diff --git a/HYBRID_SEARCH_TESTING_BRAINSTORM_2026-07-18.md b/HYBRID_SEARCH_TESTING_BRAINSTORM_2026-07-18.md new file mode 100644 index 00000000000..d3f8d71effe --- /dev/null +++ b/HYBRID_SEARCH_TESTING_BRAINSTORM_2026-07-18.md @@ -0,0 +1,155 @@ +# Hybrid-Search Testing Deep Dive — sp/hybrid-search + +**Date:** 2026-07-18 · **HEAD:** `1ab735ee` (clean) · **Mode:** god (time-boxed) · +**Contract:** correctness integration tests first, then performance attack; hybrid+BM25 scope; docker allowed; implement in-session. + +## Method + +89 blind test-attack questions (GPT-5 ×15, Gemini ×15, MiniMax ×12, 4 repo-grounded Claude +finders ×47) → merged/deduped to 35 (20 P0) → each premise verified in code by 9 grounding +agents (file:line evidence) → finalists falsified via opencode + live-cluster probes. +Artifacts: `scratchpad/bdd/{bank,verdicts,groups}.json`. + +**Verdicts:** 17 CONFIRMED-VULNERABLE · 8 PLAUSIBLE (test discriminates) · 9 SAFE (cheap +regression pins) · 1 premise-false. + +## Empirically confirmed bug (live cluster, before any new test code) + +**Q1 — `bm25(...) + @filter + first:N` returns the wrong page.** +`query/query.go:1031` pushes `First` to the worker only when `len(sg.Filters)==0`; with a +filter the worker scores everything and the coordinator slices the uid-sorted matrix +(`query/query.go:2561`) — the page is the N **lowest-UID** matches, not the N top-scored. +Probe on the test corpus: score order `0x1f7(0.659), 0x1f5(0.375)`; filtered `first:2` +returned `{0x1f5, 0x1f6}` — the top document is silently dropped. + +## Top-5 portfolio (ranked; prerequisites respected) + +### 1. Silent-wrong-page suite — pagination/ordering integrity (P0, existing compose) +Q1 (CONFIRMED — static + empirical probe + opencode, triple-verified), Q3 (tie storms/page +partition/determinism), Q12 (part 1 CONFIRMED: `orderdesc: val(b)` on `uid(b,v)` silently +deletes uids unscored in `b`, `query/query.go:2777`; part 2 — multi-key var sort fallthrough +— REFUTED: parser rejects it at `dql/parser.go:3099`). +*Why now:* the exact silent-wrong-result class the branch cannot ship with. +*Acceptance:* filtered/paginated pages equal score-ordered oracle; page walk partitions +the corpus; vector-only uids survive ordering (or explicit error). +*First action:* `TestBM25FilterThenFirstPreservesScoreOrder` (fails today → drives fix). + +### 2. Fusion/hybrid contract battery (P0, compose + unit) +Q5 (NaN/Inf query vector e2e), Q6 (hybrid topk candidate completeness vs exhaustive +oracle), Q7 (fuse-block modifiers silently ignored — parser must reject like hybrid does), +Q8 (channel validity: uid-var channels, empty channels, count-var phantom uids), Q9 (fusion +math pins: window-dependent normalization), Q29 (`fuse(a,a)` duplicate channel), Q30 +(fuse-of-fuse, shared channels), Q11 (directive matrix: @cascade/@normalize/groupby). +*Acceptance:* every silent-degrade path either errors or matches documented semantics. + +### 3. Score-binding oracle pins (P0/P1 regression, compose) +Q2/Q4 verified SAFE by code-read — pin them: per-uid `val(s)` equals unfiltered oracle +after var-bound `@filter`+pagination+two-path traversal, for bm25 AND similar_to; metric +math exact (`1/(1+d)` euclidean, raw cosine/dot). Q28 degenerate similar_to inputs +(k=0, dim mismatch, NaN vector components). +*Why:* the snapshot fix (`250b79be`) has zero e2e regression coverage; a refactor could +silently reintroduce misbinding. + +### 4. Stats-integrity + lifecycle suite (P0/P1, compose then LocalCluster) +Compose-feasible now: Q13 (idempotent re-SET churn), Q15 (intra-txn parallel RMW), Q25 +(alter cycles), Q27 (lang-tagged delete asymmetry — CONFIRMED hole). +LocalCluster wave (L-effort): Q16 (rebuild under live writes — falsification corrected the +mechanism: live writes between DropPrefix and the absolute flush RMW from an EMPTY bucket +and shadow the rebuild's total via MVCC → **undercount**; test assertion unchanged), +Q17/Q19 (multi-group parity, replica divergence), Q23 (bulk/live parity), Q22 (HNSW MVCC +lifecycle), Q24 (backup/restore). Q20 (crash-mid-rebuild forever-empty) was **REFUTED** by +opencode falsification — schema commits only after BuildIndexes and Raft WAL replay re-runs +the rebuild — downgraded to a Wave-2 crash-recovery sanity check. +*Acceptance:* stats equal recomputed oracle after every lifecycle transition; crash never +leaves silent-empty index. + +### 5. Performance attack (P2, after correctness) +Q31 WAND effectiveness (Zipfian corpora; postings-decoded counters; first:k vs exhaustive +crossover), Q33 ingest contention (writers × bucket collisions; CONFIRMED: batch >32 +docs/txn guarantees intra-txn conflicts), Q32 fusion overhead (channel size sweep), Q35 +per-query floors (32-bucket stats read), Q34 scored-HNSW overhead pin (until `WantScores` +lands). Note: local image is built WITHOUT jemalloc — absolute numbers shift, relative +comparisons valid. + +## Not eligible / deferred +- Q14 premise-false (WAND bounds are computed per-query from live postings — no stale-bound + mechanism); cheap differential pin only. +- Q18 SAFE (channel-group failure propagates as whole-query error — pin later). +- Q21/Q26 SAFE (snapshot-consistent bucket reads; namespace-prefixed keys) — unit pins. +- #8 scored-path overhead fix itself remains blocked on protoc (separate item). + +## Harness notes (established this session) +- `dgraph/dgraph:local` built via pure-Go cross-compile (`CGO_ENABLED=0 GOOS=linux`), + bypassing sudo-blocked jemalloc; binary smoke-tested. +- Cluster: `COMPOSE_COMPATIBILITY=true LINUX_GOBIN= docker compose + -p dgraph -f dgraph/docker-compose.yml up -d` (v1 underscore names required by + `testutil.getContainer`). +- Run: `TEST_DOCKER_PREFIX=dgraph go test -tags integration ./query/ -run -count=1`. +- Multi-group/kill: `dgraphtest.LocalCluster` (`WithNumAlphas`, `KillAlpha`) — feasible. + +## Wave-1 outcome (implemented this session, commit 0b806236e) + +~30 tests added across query/query_bm25_test.go, query/query_hybrid_test.go, +dql/fuse_parser_test.go, query/fuse_test.go, worker/bm25wand_test.go. Full +bm25+hybrid+fuse integration battery green against the live compose cluster. + +Six confirmed bugs found and fixed, each gated by a test: +1. **Q1** ranker root + @filter + first:N returned the lowest-uid page, not the + top-scored (applyPagination now pages by score for ranker roots). +2. **Q7/Q29** fuse() modifiers and duplicate channels silently ignored → parse errors. +3. **Q5** NaN/Inf: dropped at score snapshot; non-finite query vectors rejected. +4. **Q28** similar_to k<=0 panicked (OOM'd) the alpha → rejected before HNSW. +5. **Q8** uid-var fuse channel silently empty — root cause: ShardedMap.IsEmpty() + only detects nil (fresh map has 30 empty shards) → rejected with clear error. +6. **Q27** @lang + @index(bm25) made tagged values silently unsearchable + (lang-qualified keys; fulltext-parity) → schema combination rejected. + +Predictions that did NOT survive testing: count-var sentinel phantom uid (passes), +groupby crash unreachable via DQL (parser rejects; defensive guard added anyway), +Q20 crash-mid-rebuild (refuted by falsification pass). + +## Wave-2 outcome (commit 53caacf20) + +New `systest/bm25lifecycle` package (LocalCluster harness, `integration2` tag), all +green by default; run gated acceptance tests with `BM25_KNOWN_FAILING=1`. + +- **Fixed: intra-proposal stats race** (Q15/Q19 family) — ≥512-edge proposals apply in + parallel goroutines that never share a data key but DO share uid%32 stats buckets; + the unserialized RMW dropped deltas (heavy-txn oracle diverged). `Txn.bm25StatsMu` + serializes the stats RMW; empirically verified fixed. +- **Confirmed + gated: Q16 rebuild-vs-live undercount** — live write between DropPrefix + and the rebuild's absolute flush RMWs from an empty bucket and shadows the flush via + MVCC (uid 0x5b: 0.0062 vs closed-form 4.137). Fix design: Zero-mediated conflicting + flush or delta-merged stats. Gated test is the acceptance gate. +- **Confirmed + gated: Q23 bulk duplicate triples** — mapper counts stats per nquad, + reduce dedupes postings (implied docCount 1050 vs 1000). Fix: once-per-(pred,uid) + stats at reduce. Clean-parity (dedup'd data) asserted green by default. +- **Empirically SAFE**: Q17 multi-group parity (bm25 bit-parity; HNSW cross-topology + divergence bounded to boundary ranks — approximate-index contract documented in the + test), Q19 replica catch-up + leader failover, Q20 crash-mid-rebuild recovery, + Q22 HNSW lifecycle (no ghosts / stale scores / restart divergence). + +## Wave-3 outcome (benchmarks) + +- **WAND effectiveness** (`worker/bm25wand_bench_test.go`): topK=10 multi-term ≈ 64× + exhaustive (774/60k docs scored); single high-df term only ~11×; **topK≥100 + collapses to exhaustive**; deep offset (first:10, offset:10000 → heap 10010) ≈ 43× + slower than offset:0 — degrades to exhaustive by design. BMW == WAND on interleaved + corpora (block-max bounds collapse) — BMW currently buys nothing. +- **Fusion overhead** (`query/fuse_bench_test.go`): RRF 3.5-4× linear; `channelRanks` + sort+alloc is ~70% of 8-channel RRF at 100k/channel (581ms, 75MB); `topk` does NOT + bound work (truncation after full sort). Optimization targets: rank computation + without full map materialization; pre-fusion channel caps. +- **Stats read floor** (`posting/bm25_bench_test.go`): every bm25 query pays ~33µs / + 1100 allocs reading 32 buckets uncached (~9× a warm re-read) — per-predicate stats + snapshot cache is the obvious lever. Codec: decode 5.5ns/0 allocs (healthy). +- Q34 scored-HNSW overhead: deferred (needs populated index harness; the `WantScores` + fix is protoc-blocked anyway). + +Note: benchmarks run on a jemalloc-free build; absolute numbers shift with jemalloc, +relative comparisons hold. + +## Evidence that would change this ranking +- If Q1's fix is trivial and unblocks pushdown+rescore, suite 1 shrinks to regression pins. +- If opencode falsification refutes Q16/Q20 mechanics, the LocalCluster wave drops to P2. +- Telemetry on real query mixes (unverifiable-from-repo) would re-weight the perf attack. diff --git a/dgraph/cmd/bulk/mapper.go b/dgraph/cmd/bulk/mapper.go index 1a7299ab2cf..ba99d732567 100644 --- a/dgraph/cmd/bulk/mapper.go +++ b/dgraph/cmd/bulk/mapper.go @@ -43,6 +43,7 @@ var ( type mapper struct { *state shards []shardState // shard is based on predicate + } type shardState struct { @@ -295,8 +296,11 @@ func (m *mapper) addMapEntry(key []byte, p *pb.Posting, shard int) { atomic.AddInt64(&m.prog.mapEdgeCount, 1) uid := p.Uid - if p.PostingType != pb.Posting_REF || len(p.Facets) > 0 { - // Keep p + if p.PostingType != pb.Posting_REF || len(p.Facets) > 0 || len(p.Value) > 0 { + // Keep p. A REF posting that carries a Value (e.g. a BM25 term posting packing + // term frequency and document length) must retain that payload — mirroring the + // len(p.Value) > 0 retention clause in List.encode — or it would be reduced to a + // bare UID and the value silently lost. } else { // We only needed the UID. p = nil @@ -456,11 +460,21 @@ func (m *mapper) addIndexMapEntries(nq dql.NQuad, de *pb.DirectedEdge) { // doing edge postings. So okay to be fatal. x.Check(err) + attr := x.NamespaceAttr(nq.Namespace, nq.Predicate) + + // BM25 postings pack (term frequency, document length) into each posting's value + // and require corpus statistics; the generic token path would write bare, + // valueless postings and no stats, leaving bulk-loaded data unsearchable. Handle + // it separately. + if _, isBM25 := toker.(tok.BM25Tokenizer); isBM25 { + m.addBM25IndexMapEntries(attr, nq.Lang, de, schemaVal) + continue + } + // Extract tokens. toks, err := tok.BuildTokens(schemaVal.Value, tok.GetTokenizerForLang(toker, nq.Lang)) x.Check(err) - attr := x.NamespaceAttr(nq.Namespace, nq.Predicate) // Store index posting. for _, t := range toks { m.addMapEntry( @@ -474,3 +488,54 @@ func (m *mapper) addIndexMapEntries(nq dql.NQuad, de *pb.DirectedEdge) { } } } + +// addBM25IndexMapEntries writes the BM25 term postings for one document — one posting +// per distinct term, packing (term frequency, document length) into the value exactly +// as the live index path does — and accumulates the document's contribution to this +// mapper's corpus statistics. The accumulated stats are flushed once per bucket after +// the map phase (loader.flushBM25Stats), because corpus statistics must be summed +// across documents rather than unioned like postings. +func (m *mapper) addBM25IndexMapEntries(attr string, lang string, de *pb.DirectedEdge, + schemaVal types.Val) { + termFreqs, docLen, err := tok.BM25Tokenizer{}.TokensWithFrequency(schemaVal.Value, lang) + x.Check(err) + if docLen == 0 { + // Document tokenizes to zero terms (e.g. all stopwords); it contributes no + // postings and no corpus statistics. + return + } + + shard := m.state.shards.shardFor(attr) + uid := de.GetEntity() + for term, tf := range termFreqs { + encodedTerm := string([]byte{tok.IdentBM25}) + term + m.addMapEntry( + x.IndexKey(attr, encodedTerm), + &pb.Posting{ + Uid: uid, + PostingType: pb.Posting_REF, + ValType: pb.Posting_BINARY, + Value: posting.EncodeBM25Value(tf, docLen), + }, + shard, + ) + } + + // Emit the document's corpus-statistics contribution as a PER-DOCUMENT posting + // on the final stats bucket key: (docCount=1, totalTerms=docLen). Duplicate + // nquads for the same (attr, uid) collapse in the reducer's same-uid dedup — + // exactly like the term postings above — so duplicates cannot inflate the + // stats (mapper-side pre-aggregation counted once per nquad and did). The + // reducer folds each bucket's per-doc postings into the single aggregate + // posting the live/rebuild paths store (see reduce.go appendToList). + m.addMapEntry( + x.BM25StatsKey(attr, int(uid%posting.NumBM25StatsBuckets)), + &pb.Posting{ + Uid: uid, + PostingType: pb.Posting_REF, + ValType: pb.Posting_BINARY, + Value: posting.EncodeBM25Stats(1, uint64(docLen)), + }, + shard, + ) +} diff --git a/dgraph/cmd/bulk/mapper_test.go b/dgraph/cmd/bulk/mapper_test.go new file mode 100644 index 00000000000..3598252000f --- /dev/null +++ b/dgraph/cmd/bulk/mapper_test.go @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package bulk + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/posting" +) + +// TestBM25PerDocStatsFold verifies the reducer-side fold semantics that make bulk +// BM25 stats exact: the mapper emits one (docCount=1, totalTerms=docLen) posting +// PER DOCUMENT on the bucket's stats key, the reducer dedupes postings with the +// same uid (exactly like term postings — so duplicate nquads cannot inflate the +// stats), and the deduped contributions sum into the aggregate bucket posting. +// This test models the fold: dedup by uid, then sum decoded contributions. +func TestBM25PerDocStatsFold(t *testing.T) { + type docEntry struct { + uid uint64 + docLen uint64 + } + // Simulated post-sort reducer input for one bucket key: uid-ascending with + // duplicates (the duplicate-nquad case that inflated stats before the fold). + entries := []docEntry{ + {uid: 1, docLen: 10}, + {uid: 1, docLen: 10}, // duplicate nquad — must count once + {uid: 33, docLen: 5}, + {uid: 33, docLen: 5}, // duplicate + {uid: 33, docLen: 5}, // triplicate + {uid: 65, docLen: 7}, + } + + var docCount, totalTerms uint64 + var lastUid uint64 + for _, e := range entries { + if e.uid == lastUid { + continue + } + lastUid = e.uid + // Round-trip through the wire format the mapper emits. + c, terms := posting.DecodeBM25Stats(posting.EncodeBM25Stats(1, e.docLen)) + docCount += c + totalTerms += terms + } + + require.Equal(t, uint64(3), docCount, "duplicates must not inflate docCount") + require.Equal(t, uint64(22), totalTerms, "duplicates must not inflate totalTerms") +} diff --git a/dgraph/cmd/bulk/reduce.go b/dgraph/cmd/bulk/reduce.go index d8cb5385514..eee1083f97d 100644 --- a/dgraph/cmd/bulk/reduce.go +++ b/dgraph/cmd/bulk/reduce.go @@ -811,6 +811,54 @@ func (r *reducer) toList(req *encodeRequest, vi *vectorIndexer) { enc := codec.Encoder{BlockSize: 256, Alloc: alloc} var lastUid uint64 var slice []byte + + // BM25 corpus-statistics bucket: the mapper emits one (docCount=1, + // totalTerms=docLen) posting PER DOCUMENT on the bucket key, so duplicate + // nquads for the same (attr, uid) collapse in the same-uid dedup below — + // they cannot inflate the stats the way mapper-side pre-aggregation did + // (it counted once per nquad). Fold the deduped per-doc postings into the + // single aggregate value posting the live/rebuild paths store. + if pk.IsBM25Stats() { + var docCount, totalTerms uint64 + var p pb.Posting + lastUid = 0 + next := start + for next >= 0 && (next < end || end == -1) { + slice, next = cbuf.Slice(next) + me := MapEntry(slice) + uid := me.Uid() + if uid == lastUid { + continue + } + lastUid = uid + if pbuf := me.Plist(); len(pbuf) > 0 { + p.Reset() + x.Check(proto.Unmarshal(pbuf, &p)) + c, terms := posting.DecodeBM25Stats(p.Value) + docCount += c + totalTerms += terms + } + } + if docCount == 0 && totalTerms == 0 { + return + } + enc.Add(math.MaxUint64) + pl.Postings = append(pl.Postings, &pb.Posting{ + Uid: math.MaxUint64, + PostingType: pb.Posting_VALUE, + ValType: pb.Posting_BINARY, + Value: posting.EncodeBM25Stats(docCount, totalTerms), + }) + pl.Pack = enc.Done() + kv := posting.MarshalPostingList(pl, nil) + kv.Key = y.Copy(currentKey) + kv.Version = writeVersionTs + kv.StreamId = r.streamIdFor(pk.Attr) + badger.KVToBuffer(kv, kvBuf) + pl.Reset() + return + } + next := start for next >= 0 && (next < end || end == -1) { slice, next = cbuf.Slice(next) diff --git a/docs/superpowers/specs/2026-06-03-hybrid-search-fusion-design.md b/docs/superpowers/specs/2026-06-03-hybrid-search-fusion-design.md new file mode 100644 index 00000000000..fa698e17272 --- /dev/null +++ b/docs/superpowers/specs/2026-06-03-hybrid-search-fusion-design.md @@ -0,0 +1,255 @@ +# Native Hybrid Search with Score Fusion — Design Spec + +**Branch:** `sp/hybrid-search` (off `sp/bm25`) **Date:** 2026-06-03 **Status:** Approved (Option C) + +## 1. Problem + +Dgraph can rank text (BM25, on `sp/bm25`) and find nearest vectors (HNSW `similar_to`), but cannot +**combine** them. A consumer wanting hybrid retrieval (the standard RAG pattern: dense vector + +sparse/keyword, fused into one ranked list) must issue separate DQL queries and fuse the results in +application code. + +Concretely, the reference consumer (modelhub, a GraphRAG product running on Dgraph) issues **three** +separate DQL queries per search — (1) `similar_to` vector, (2) BM25-ish term, (3) entity-anchored +term — and fuses them in Python with Reciprocal Rank Fusion (`score = Σ 1/(k+rank)`, k=60), +sometimes with a linear combination (`α·vec + (1-α)·text`, scores max-normalized first). + +This spec brings that fusion **natively into a single DQL query**. + +## 2. Approach (Option C) + +Two pieces, agreed after consulting GPT-5 and Gemini (both recommended the N-way primitive as the +core; Option C adds a thin convenience wrapper): + +1. **`fuse()`** — an N-way fusion combinator over already-scored DQL **value variables**. The + general primitive. Handles any number of channels and any scored signal, not just bm25/vector. +2. **`hybrid()`** — convenience sugar for the common 2-channel bm25+vector case, expanded at + query-rewrite time into two channel blocks + a `fuse()`. No new executor path. + +A prerequisite makes `fuse()` useful with vectors: + +1. **Surface `similar_to` similarity scores as a value variable** (today it returns only uids). + Required so vector results can be a fusion channel. + +### Why this fits Dgraph's architecture + +- DQL already has **value variables** carrying uid→score maps (`var(func: bm25(...))` binds per-doc + scores into a `varValue{Uids, Vals}` — see `query/query.go` `populateUidValVar`, the `bm25` case). +- `ProcessQuery` already has a **variable-dependency scheduler** (`canExecute` over + `QueryVars.Needs`/`.Defines`): a block runs only once the variables it needs are populated. A + `fuse()` block that _needs_ its channel vars and _defines_ the fused var is scheduled + automatically after its inputs — no new sequencing logic. +- Fusion is a **coordinator-side** operation over resolved variables (like `math()` and `uid(a,b)`), + which is correct under predicate sharding: each channel may be computed on a different Raft group; + their value variables are already merged to the coordinator before fusion runs. + +## 3. DQL Surface + +### 3.1 `fuse()` — N-way primitive + +```text +v as var(func: bm25(text, "quick brown fox")) +e as var(func: similar_to(emb, 100, $queryVec)) + +f as var(func: fuse(v, e, method: "rrf", k: 60)) + +{ + result(func: uid(f), orderdesc: val(f), first: 10) { + uid + val(f) + } +} +``` + +- **Positional args**: two or more value-variable references (the channels). These become the + block's `NeedsVar`. +- **Named options** (parsed like `similar_to`'s `ef:`/`distance_threshold:`): + - `method`: `"rrf"` (default) or `"linear"`. + - `k`: RRF rank constant (default `60`). Ignored for linear. + - `weights`: comma-separated per-channel weights for linear, aligned positionally with the channel + args (e.g. `weights: "0.3,0.7"`). Default: `1.0` each. Ignored for RRF. + - `normalize`: linear-only score normalization, `"max"` (default) or `"none"`. + - `topk`: optional cap on the number of fused results emitted (default: unbounded; downstream + `first/offset` still applies). Bounds coordinator work. +- **Output**: binds `f` as a value variable — both the **union** uid set (`uid(f)`) and the + uid→fusedScore map (`val(f)`, `orderdesc: val(f)`), exactly like the bm25 ranker var. + +### 3.2 `hybrid()` — 2-channel sugar + +```text +f as var(func: hybrid(text, "quick brown fox", emb, $queryVec, 100, method: "rrf", k: 60)) + +{ result(func: uid(f), orderdesc: val(f), first: 10) { uid val(f) } } +``` + +Positional: `textPredicate, "queryText", vectorPredicate, $queryVec, topk`. Named options as for +`fuse()`. Rewritten at query-build time to: + +```text +__hybrid_0_ as var(func: bm25(text, "quick brown fox")) +__hybrid_1_ as var(func: similar_to(emb, 100, $queryVec)) +f as var(func: fuse(__hybrid_0_, __hybrid_1_, method: "rrf", k: 60)) +``` + +The synthetic var names are unique per hybrid block. After rewrite there is no distinct `hybrid` +execution path — it is purely a parser/builder transformation. + +## 4. Fusion Semantics (the core, `query/fuse.go`) + +Pure function over channels, fully unit-testable, no I/O: + +```go +type fuseChannel struct { + scores map[uint64]float64 // uid -> raw channel score (higher = better) + weight float64 // linear weight (default 1.0) +} + +func fuseRRF(channels []fuseChannel, k float64, topk int) []scoredUid +func fuseLinear(channels []fuseChannel, normalize string, topk int) []scoredUid +``` + +### Outer-join / union semantics (the key correctness point) + +Both models independently flagged this. Fusion is a **set union** of candidate uids across channels, +NOT an intersection. For a uid present in some channels but not others: + +- **RRF**: only ranked channels contribute. `fused[u] = Σ_{c: u∈c} 1/(k + rank_c(u))`. A channel + that doesn't contain `u` adds nothing (equivalent to rank = ∞). +- **Linear**: missing channel contributes `0`. `fused[u] = Σ_c weight_c · norm_c(score_c(u))`, where + `norm_c(missing) = 0`. + +Standard DQL `math()` across var blocks aligns/intersects on uid and would drop or NaN-poison such +uids — `fuse()` must not. + +### RRF ranks + +Each channel is independently sorted by raw score **descending**, tie-broken by uid **ascending**, +to assign 1-based ranks. (Deterministic; matches the bm25/HNSW `sorted()` tie-break already in the +codebase.) + +### Linear normalization + +`"max"` (default): divide each channel's scores by that channel's max (`|max|`, guard against 0 → +channel contributes 0). Brings heterogeneous score scales (BM25 ∈ [0,∞), cosine ∈ [-1,1]) onto a +comparable range before weighting. `"none"`: use raw scores (caller asserts they're comparable). + +### Output ordering + +Fused results sorted by fused score descending, tie-broken by uid ascending. If `topk > 0`, truncate +to `topk`. The output is emitted to the value variable as the union uid set (ascending, per the +query pipeline contract) + positionally-aligned fused scores — identical shape to the bm25 var +binding. + +## 5. Vector score surfacing (prerequisite) + +`similar_to` currently emits only `UidMatrix` (worker/task.go ~L417). Change: + +- Extend `index.SearchPathResult` with `Distances []float64` parallel to `Neighbors`, populated from + the search-layer heap (which already holds metric-domain distances, `n.value`) in + `addFinalNeighbors`. +- In the `similarToFn` worker path, when the function is bound to a value variable, use the + path-returning search to obtain distances and emit a `ValueMatrix` of **similarity** scores + (higher = better): + - **cosine**: cosine similarity in [-1,1], as-is. + - **dotproduct**: dot product, as-is. + - **euclidean**: `1/(1 + d)` where `d` is the (non-squared) Euclidean distance, mapping to (0,1] + so higher = better and linear normalization is well-behaved. +- Bind in `populateUidValVar` exactly like the `bm25` case (reuse/generalize that branch). When + `similar_to` is **not** bound to a var, behavior is unchanged (uids only) — zero overhead for + existing queries. + +Orientation contract: **all rankers surface higher-is-better scores**, so RRF and linear fusion +compose without per-channel sign handling. + +## 6. Execution flow + +1. Parse: `fuse`/`hybrid` recognized as functions in `dql/parser.go`; channel args → `NeedsVar`, + named options stored on the function. `hybrid` expands to channel blocks + `fuse`. The block + `Defines` its output var, `Needs` its channel vars. +2. Schedule: `ProcessQuery`'s `canExecute` runs channel blocks first; the `fuse` block becomes + runnable once all channel vars are in `req.Vars`. +3. Compute: the `fuse` block is **coordinator-only** — like the existing + `similar_to`-empty/`IsEmpty` cases, it is **not** dispatched to a worker. Fusion is computed in + `populateVarMap` (new `fuse` case) reading channel `varValue`s from `doneVars`, producing the + fused `varValue`. +4. Consume: downstream `uid(f)` / `orderdesc: val(f)` / `first`/`offset` work via the existing + value-variable machinery. + +## 7. Validation & errors + +- ≥1 channel var required (2+ to be meaningful; 1 is allowed and passes through). +- Unknown `method` → error. `k <= 0` → error. `weights` count must match channel count when provided + → error. Malformed `weights` floats → error. +- Empty channels (no matches) are valid and contribute nothing. +- A channel var that is a **uid variable without scores** (no `Vals`): for RRF, rank by the var's + intrinsic order if any, else treat as unscored → error with a clear message ("fuse channel %q has + no scores; use a ranker like bm25/similar_to"). MVP: require scored channels. + +## 8. Testing + +**Unit (`query/fuse_test.go`)** — pure fusion core: + +- RRF: known ranks → known `Σ 1/(k+rank)`; default k=60; custom k. +- Linear: max-normalize, weights, `normalize:none`. +- Union semantics: uid in 1 of N channels; disjoint channels; full overlap. +- Ties (equal scores → uid-ascending rank); empty channels; single channel. +- `topk` truncation; determinism. + +**Worker (`worker/`)** — vector score surfacing: + +- `similar_to` bound to a var emits similarity scores with correct orientation per metric; unbound + `similar_to` unchanged. + +**Integration (systest/DQL)** — end-to-end: + +- `fuse()` RRF over bm25 + vector; ordering matches hand-computed RRF. +- `fuse()` linear with weights. +- 3-channel fusion (modelhub's shape). +- Pagination (`first`/`offset`) on the consuming block. +- Missing-uid union correctness. +- `hybrid()` produces results identical to the equivalent explicit `fuse()`. +- Error paths (unknown method, bad weights, unscored channel). + +## 9. Out of scope (future) + +- Filter pushdown into HNSW (pre-filtered ANN) — separate gap. +- Worker-side fusion pushdown / `topk` propagation into channel funcs. +- Additional methods (weighted-RRF, ISR, distribution-based fusion). +- Reranking primitives. + +## 9a. Adversarial review outcomes (GPT-5 + Gemini) + +Both models deep-reviewed the diff. Findings triaged and resolved: + +- **Behavior preservation for `similar_to` (High, both):** the worker no longer routes plain + (no-option) vector queries through the options path. New `SearchScored` / `SearchWithUidScored` + mirror `Search` / `SearchWithUid` exactly and just also return scores, so existing queries' + neighbor selection is unchanged; the `*Options*` scored variants are used only when + ef/distance-threshold is given. +- **NaN/Inf safety (both):** `scoresFromVar` drops non-finite scores and `channelMaxAbs` ignores + them, so a pathological score can never break the sort comparator's strict-weak-ordering or poison + a linear sum. +- **"Ascending Uids destroys ranking" (both, flagged Critical):** not a bug — this follows the same + value-variable contract as bm25 (Uids is the unordered set for `uid(var)`; ranked order is + recovered via `orderdesc: val(var)`; `topk` selection happens before the ascending sort). + Documented in `computeFuse`. +- **Undefined channel var (GPT-5 Critical "stall"):** not a stall — `checkDependency` rejects it at + parse time ("variables used but not defined"). Regression test added. +- **Synthetic var collision (both, Low):** `__hybrid` is now a reserved prefix; user vars using it + are rejected with a clear message. +- **`@filter(similar_to(...))` + ValueMatrix (both, Med):** follows the proven bm25 precedent (bm25 + emits ValueMatrix and is used in filters); to be confirmed by the vector integration suite in CI. +- **Coordinator GC churn in `channelRanks` (Gemini, perf):** acknowledged; acceptable at expected + channel sizes, noted as future optimization (buffer pooling) if needed. + +## 10. Files touched + +| Area | File(s) | +| --------------------- | ------------------------------------------------------------------------------------------------------------------------------- | +| Fusion core (new) | `query/fuse.go`, `query/fuse_test.go` | +| Parser | `dql/parser.go` (recognize `fuse`/`hybrid`, args+opts), `dql/state.go` if needed | +| hybrid rewrite | `query/query.go` (ToSubGraph/build) or `dql` transform | +| Var binding | `query/query.go` `populateUidValVar` (fuse case; generalize bm25 case) | +| Scheduler skip-worker | `query/query.go` `ProcessQuery` (fuse → no worker dispatch) | +| Vector scores | `tok/index/search_path.go` (`Distances`), `tok/hnsw/persistent_hnsw.go` (populate), `worker/task.go` (`similar_to` ValueMatrix) | +| Docs | DQL docs for `fuse`/`hybrid` | diff --git a/dql/fuse_parser_test.go b/dql/fuse_parser_test.go new file mode 100644 index 00000000000..e42ad950f68 --- /dev/null +++ b/dql/fuse_parser_test.go @@ -0,0 +1,537 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package dql + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// findBlock returns the query block whose result variable is varName. +func findVarBlock(t *testing.T, res Result, varName string) *GraphQuery { + for _, q := range res.Query { + if q.Var == varName { + return q + } + } + t.Fatalf("no block defines var %q", varName) + return nil +} + +func TestParseFuse_ChannelsAndOptions(t *testing.T) { + query := ` + { + v as var(func: bm25(text, "quick brown fox")) + e as var(func: similar_to(emb, 100, "[0.1, 0.2]")) + f as var(func: fuse(v, e, method: "rrf", k: 60)) + result(func: uid(f), orderdesc: val(f)) { + uid + } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + + fb := findVarBlock(t, res, "f") + require.NotNil(t, fb.Func) + require.Equal(t, "fuse", fb.Func.Name) + + // Channels are captured as value-variable NeedsVar in order. + require.Len(t, fb.Func.NeedsVar, 2) + require.Equal(t, "v", fb.Func.NeedsVar[0].Name) + require.Equal(t, ValueVar, fb.Func.NeedsVar[0].Typ) + require.Equal(t, "e", fb.Func.NeedsVar[1].Name) + + // Named options are captured as [key, value] arg pairs. + args := map[string]string{} + for i := 0; i+1 < len(fb.Func.Args); i += 2 { + args[fb.Func.Args[i].Value] = fb.Func.Args[i+1].Value + } + require.Equal(t, "rrf", args["method"]) + require.Equal(t, "60", args["k"]) +} + +func TestParseFuse_LinearWeights(t *testing.T) { + query := ` + { + v as var(func: bm25(text, "fox")) + e as var(func: bm25(title, "fox")) + f as var(func: fuse(v, e, method: "linear", weights: "0.3,0.7", normalize: "max")) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + fb := findVarBlock(t, res, "f") + require.Len(t, fb.Func.NeedsVar, 2) + args := map[string]string{} + for i := 0; i+1 < len(fb.Func.Args); i += 2 { + args[fb.Func.Args[i].Value] = fb.Func.Args[i+1].Value + } + require.Equal(t, "linear", args["method"]) + require.Equal(t, "0.3,0.7", args["weights"]) + require.Equal(t, "max", args["normalize"]) +} + +func TestParseFuse_ThreeChannels(t *testing.T) { + query := ` + { + a as var(func: bm25(text, "fox")) + b as var(func: bm25(title, "fox")) + c as var(func: bm25(body, "fox")) + f as var(func: fuse(a, b, c)) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + fb := findVarBlock(t, res, "f") + require.Len(t, fb.Func.NeedsVar, 3) +} + +func TestParseFuse_UnknownOption(t *testing.T) { + query := ` + { + v as var(func: bm25(text, "fox")) + f as var(func: fuse(v, bogus: "x")) + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "Unknown option") +} + +func TestParseFuse_DuplicateOption(t *testing.T) { + query := ` + { + v as var(func: bm25(text, "fox")) + f as var(func: fuse(v, k: 10, k: 20)) + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "Duplicate key") +} + +func TestParseFuse_NoChannels(t *testing.T) { + query := ` + { + f as var(func: fuse(method: "rrf")) + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "at least one value variable") +} + +func TestParseHybrid_ExpandsToThreeBlocks(t *testing.T) { + query := ` + { + f as var(func: hybrid(description, "quick brown fox", emb, "[0.1, 0.2]", topk: 50, method: "rrf", k: 60)) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + + // The hybrid block is replaced by bm25 + similar_to + fuse (plus the result block). + var bm25Block, simBlock, fuseBlock *GraphQuery + for _, q := range res.Query { + if q.Func == nil { + continue + } + switch q.Func.Name { + case "bm25": + bm25Block = q + case "similar_to": + simBlock = q + case "fuse": + fuseBlock = q + case "hybrid": + t.Fatal("hybrid block should have been rewritten away") + } + } + require.NotNil(t, bm25Block, "bm25 channel block must exist") + require.NotNil(t, simBlock, "similar_to channel block must exist") + require.NotNil(t, fuseBlock, "fuse block must exist") + + // bm25 channel: predicate + query text. + require.Equal(t, "description", bm25Block.Func.Attr) + require.Equal(t, "quick brown fox", bm25Block.Func.Args[0].Value) + + // similar_to channel: predicate + topk + vector. + require.Equal(t, "emb", simBlock.Func.Attr) + require.Equal(t, "50", simBlock.Func.Args[0].Value) + + // fuse block keeps the original variable name and the fuse options. + require.Equal(t, "f", fuseBlock.Var) + require.Len(t, fuseBlock.Func.NeedsVar, 2) + args := map[string]string{} + for i := 0; i+1 < len(fuseBlock.Func.Args); i += 2 { + args[fuseBlock.Func.Args[i].Value] = fuseBlock.Func.Args[i+1].Value + } + require.Equal(t, "rrf", args["method"]) + require.Equal(t, "60", args["k"]) + // topk is consumed by similar_to, not forwarded to fuse. + require.NotContains(t, args, "topk") +} + +func TestParseHybrid_BoundsBM25Channel(t *testing.T) { + // The generated bm25 channel must be bounded to topk so a broad text query does + // not score the whole corpus before fusion. + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]", topk: 25)) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + for _, q := range res.Query { + if q.Func != nil && q.Func.Name == "bm25" { + require.Equal(t, "25", q.Args["first"], "bm25 channel should be capped at topk") + } + } +} + +func TestParseHybrid_MalformedOptions(t *testing.T) { + // A trailing option key without a value must be rejected, not silently dropped. + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]", method)) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) +} + +func TestParseHybrid_DefaultTopK(t *testing.T) { + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]")) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + for _, q := range res.Query { + if q.Func != nil && q.Func.Name == "similar_to" { + require.Equal(t, "100", q.Func.Args[0].Value, "default topk should be 100") + } + } +} + +func TestParseFuse_UndefinedChannelVarErrors(t *testing.T) { + // A fuse channel referencing a variable that no block defines must be rejected + // at parse time (not silently stall the scheduler). + query := ` + { + v as var(func: bm25(text, "fox")) + f as var(func: fuse(v, ghost, method: "rrf")) + result(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "not defined") +} + +func TestParseHybrid_RequiresVar(t *testing.T) { + query := ` + { + result(func: hybrid(description, "fox", emb, "[0.1]")) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "must be assigned to a variable") +} + +// fuseArgPairs flattens a fuse Function's [key, value, key, value, ...] Args +// into a map for assertions. +func fuseArgPairs(fn *Function) map[string]string { + args := map[string]string{} + for i := 0; i+1 < len(fn.Args); i += 2 { + args[fn.Args[i].Value] = fn.Args[i+1].Value + } + return args +} + +// Regression gate (Q7): fuse blocks used to silently accept @filter/first/offset/ +// order/@cascade/children even though execution never applies them (query/query.go +// skips ProcessGraph for fuse blocks). validateFuseBlocks (dql/hybrid.go) must +// reject every such modifier, mirroring hybrid's guard. +func TestParseFuse_RejectsBlockModifiers(t *testing.T) { + cases := []struct { + name string + block string + }{ + {"filter", `f as var(func: fuse(a, b)) @filter(has(name))`}, + {"first", `f as var(func: fuse(a, b), first: 5)`}, + {"offset", `f as var(func: fuse(a, b), offset: 2)`}, + {"orderdesc", `f as var(func: fuse(a, b), orderdesc: val(a))`}, + {"cascade", `f as var(func: fuse(a, b)) @cascade`}, + {"children", `f as var(func: fuse(a, b)) { name }`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + query := ` + { + a as var(func: bm25(text, "fox")) + b as var(func: bm25(title, "fox")) + ` + tc.block + ` + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err, + "fuse block modifier %q must be rejected, not silently ignored", tc.name) + require.Contains(t, err.Error(), "not supported on a fuse() block") + }) + } +} + +// Regression gate (Q7): hybrid's modifier guard originally covered +// @filter/order/cascade/children but not pagination Args — `first`/`offset` on a +// hybrid() block parsed fine and were then silently discarded by expandHybridBlock. +// They must error like the other modifiers. +func TestParseHybrid_RejectsPagination(t *testing.T) { + for _, mod := range []string{"first: 5", "offset: 2"} { + t.Run(mod, func(t *testing.T) { + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]"), ` + mod + `) + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err, + "hybrid pagination %q must be rejected, not silently dropped", mod) + require.Contains(t, err.Error(), "not supported on a hybrid() block") + }) + } +} + +// Regression pin (Q7/Q10, dc12fbbc1): hybrid's existing modifier guard rejects +// @filter, ordering, @cascade, and child blocks with an actionable message. +func TestParseHybrid_RejectsFilterOrderCascadeChildren(t *testing.T) { + cases := []struct { + name string + block string + }{ + {"filter", `f as var(func: hybrid(description, "fox", emb, "[0.1]")) @filter(has(name))`}, + {"orderdesc", `f as var(func: hybrid(description, "fox", emb, "[0.1]"), orderdesc: name)`}, + {"cascade", `f as var(func: hybrid(description, "fox", emb, "[0.1]")) @cascade`}, + {"children", `f as var(func: hybrid(description, "fox", emb, "[0.1]")) { name }`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + query := ` + { + ` + tc.block + ` + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "not supported on a hybrid() block") + }) + } +} + +// Regression gate (Q29): fuse(a, a) used to parse with the same channel twice in +// NeedsVar, so RRF added the duplicate channel's 1/(k+rank) twice per uid. A +// duplicate channel must be rejected (or, equivalently, deduped) — never silently +// double-counted. +func TestParseFuse_DuplicateChannelVar(t *testing.T) { + query := ` + { + a as var(func: bm25(text, "fox")) + f as var(func: fuse(a, a)) + result(func: uid(f)) { uid } + }` + res, err := Parse(Request{Str: query}) + if err != nil { + // Rejection is the current resolution (validateFuseBlocks). + require.Contains(t, err.Error(), "duplicate channel") + return + } + // A documented dedup of the channel list would also be acceptable. + fb := findVarBlock(t, res, "f") + require.Len(t, fb.Func.NeedsVar, 1, + "duplicate fuse channel must be rejected or deduped, not double-counted") +} + +// Pins CURRENT deterministic parser behavior for adversarial fuse option values. +// Range/semantic validation of option values (negative k, malformed method) is an +// execution-time concern; the parser only guarantees the shapes asserted here. +func TestParseFuse_AdversarialOptionValues(t *testing.T) { + parse := func(opts string) (Result, error) { + query := ` + { + a as var(func: bm25(text, "fox")) + b as var(func: bm25(title, "fox")) + f as var(func: fuse(a, b` + opts + `)) + result(func: uid(f)) { uid } + }` + return Parse(Request{Str: query}) + } + + t.Run("negative k unquoted is a lex error", func(t *testing.T) { + _, err := parse(`, k: -1`) + require.Error(t, err) + require.Contains(t, err.Error(), "Expected value for k") + }) + + t.Run("negative k quoted parses through", func(t *testing.T) { + // The parser treats option values as opaque strings; k range checks happen + // at execution time. + res, err := parse(`, k: "-1"`) + require.NoError(t, err) + require.Equal(t, "-1", fuseArgPairs(findVarBlock(t, res, "f").Func)["k"]) + }) + + t.Run("k with no value is a lex error", func(t *testing.T) { + _, err := parse(`, k:`) + require.Error(t, err) + require.Contains(t, err.Error(), "Expected value for k") + }) + + t.Run("unquoted weights list fails var validation", func(t *testing.T) { + // weights: 0.7,0.3 lexes "0.3" as an extra bare channel name, which then + // fails loudly as an undefined variable rather than silently mis-weighting. + _, err := parse(`, method: "linear", weights: 0.7,0.3`) + require.Error(t, err) + require.Contains(t, err.Error(), "used but not defined") + }) + + t.Run("quoted weights list parses", func(t *testing.T) { + res, err := parse(`, method: "linear", weights: "0.7,0.3"`) + require.NoError(t, err) + require.Equal(t, "0.7,0.3", fuseArgPairs(findVarBlock(t, res, "f").Func)["weights"]) + }) + + t.Run("method value with colon parses through", func(t *testing.T) { + // An unknown method string is rejected at execution time, not by the parser. + res, err := parse(`, method: "rr:f"`) + require.NoError(t, err) + require.Equal(t, "rr:f", fuseArgPairs(findVarBlock(t, res, "f").Func)["method"]) + }) +} + +// Q10 hygiene pin: two hybrid() blocks in one query expand with distinct +// monotonically-indexed channel names, and each fuse block references only its +// own channels — no cross-contamination. +func TestParseHybrid_TwoBlocksUniqueNames(t *testing.T) { + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]")) + g as var(func: hybrid(title, "dog", emb2, "[0.2]", topk: 7)) + r1(func: uid(f)) { uid } + r2(func: uid(g)) { uid } + }` + res, err := Parse(Request{Str: query}) + require.NoError(t, err) + + // Every generated channel var must be unique across both expansions. + seen := map[string]bool{} + for _, q := range res.Query { + if q.Var == "" { + continue + } + require.False(t, seen[q.Var], "variable %q defined twice", q.Var) + seen[q.Var] = true + } + require.True(t, seen["__hybrid0_bm25"]) + require.True(t, seen["__hybrid0_vec"]) + require.True(t, seen["__hybrid1_bm25"]) + require.True(t, seen["__hybrid1_vec"]) + + // Each fuse block consumes exactly its own hybrid's channels. + fb := findVarBlock(t, res, "f") + require.Equal(t, "fuse", fb.Func.Name) + require.Len(t, fb.Func.NeedsVar, 2) + require.Equal(t, "__hybrid0_bm25", fb.Func.NeedsVar[0].Name) + require.Equal(t, "__hybrid0_vec", fb.Func.NeedsVar[1].Name) + + gb := findVarBlock(t, res, "g") + require.Equal(t, "fuse", gb.Func.Name) + require.Len(t, gb.Func.NeedsVar, 2) + require.Equal(t, "__hybrid1_bm25", gb.Func.NeedsVar[0].Name) + require.Equal(t, "__hybrid1_vec", gb.Func.NeedsVar[1].Name) + + // Per-block options stay per-block: the second hybrid's topk bounds only its + // own channels. + require.Equal(t, "7", findVarBlock(t, res, "__hybrid1_bm25").Args["first"]) + require.Equal(t, "7", findVarBlock(t, res, "__hybrid1_vec").Func.Args[0].Value) + require.Equal(t, "100", findVarBlock(t, res, "__hybrid0_bm25").Args["first"]) +} + +// Q10 hygiene pin: user variables using the reserved __hybrid prefix are rejected +// before the rewrite, whether declared at a root block or nested inside a child. +func TestParseHybrid_ReservedPrefixRejected(t *testing.T) { + t.Run("root", func(t *testing.T) { + query := ` + { + __hybrid0_bm25 as var(func: bm25(text, "fox")) + f as var(func: hybrid(description, "fox", emb, "[0.1]")) + r(func: uid(f, __hybrid0_bm25)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "reserved prefix") + }) + t.Run("nested child", func(t *testing.T) { + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]")) + r(func: uid(f)) { __hybrid0_vec as name } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), "reserved prefix") + }) +} + +// Q10 hygiene pin: a parameterized hybrid query text ($q) must reach the generated +// bm25 channel substituted, not as the literal "$q" (regression for dc12fbbc1). +func TestParseHybrid_DollarVarTextArg(t *testing.T) { + query := ` + query q($q: string, $vec: string) { + f as var(func: hybrid(description, $q, emb, $vec)) + r(func: uid(f)) { uid } + }` + res, err := Parse(Request{ + Str: query, + Variables: map[string]string{"$q": "lazy dog", "$vec": "[0.9, 0.8]"}, + }) + require.NoError(t, err) + + bb := findVarBlock(t, res, "__hybrid0_bm25") + require.Equal(t, "bm25", bb.Func.Name) + require.Len(t, bb.Func.Args, 1) + require.Equal(t, "lazy dog", bb.Func.Args[0].Value, + "bm25 channel must receive the substituted text, not the literal $q") + + vb := findVarBlock(t, res, "__hybrid0_vec") + require.Equal(t, "[0.9, 0.8]", vb.Func.Args[1].Value, + "similar_to channel must receive the substituted vector") +} + +// Q10 hygiene pin: non-positive topk is rejected up front (a non-positive value +// would reach HNSW as expectedNeighbors<=0 and panic the search). +func TestParseHybrid_NonPositiveTopKRejected(t *testing.T) { + cases := []struct { + name, topk, wantErr string + }{ + {"zero", `0`, "topk must be a positive integer"}, + {"negative quoted", `"-1"`, "topk must be a positive integer"}, + // Unquoted negatives never lex as an option value at all. + {"negative unquoted", `-1`, "Expected value for topk"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + query := ` + { + f as var(func: hybrid(description, "fox", emb, "[0.1]", topk: ` + tc.topk + `)) + result(func: uid(f)) { uid } + }` + _, err := Parse(Request{Str: query}) + require.Error(t, err) + require.Contains(t, err.Error(), tc.wantErr) + }) + } +} diff --git a/dql/hybrid.go b/dql/hybrid.go new file mode 100644 index 00000000000..def38ada6e6 --- /dev/null +++ b/dql/hybrid.go @@ -0,0 +1,224 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package dql + +import ( + "fmt" + "strconv" + "strings" +) + +// hybrid() is convenience sugar for the common two-channel hybrid-search case: +// BM25 text relevance fused with vector similarity. It is rewritten, before +// variable collection and dependency checking, into the equivalent explicit form: +// +// x as var(func: hybrid(textPred, "query", vecPred, $vec, topk: 100, method: "rrf", k: 60)) +// +// becomes +// +// __hybrid0_bm25 as var(func: bm25(textPred, "query"), first: 100) +// __hybrid0_vec as var(func: similar_to(vecPred, 100, $vec)) +// x as var(func: fuse(__hybrid0_bm25, __hybrid0_vec, method: "rrf", k: 60)) +// +// There is therefore no distinct hybrid execution path: it desugars entirely to +// the fuse() primitive. Positional args are textPred, "query", vecPred and the +// query vector ($var or literal); named options are topk (vector neighbor count, +// default 100) plus the fuse options method/k/weights/normalize. +const ( + hybridTopKOption = "topk" + hybridDefaultTopK = "100" + // hybridVarPrefix namespaces the synthetic channel variables a hybrid() block + // expands into. It is reserved: user variables may not start with it. + hybridVarPrefix = "__hybrid" +) + +// rewriteHybridBlocks expands every hybrid() query block in res into its three +// constituent blocks. It runs before fragment expansion, variable substitution, +// and dependency checking so the generated blocks participate normally. +func rewriteHybridBlocks(res *Result) error { + hasHybrid := false + for _, qu := range res.Query { + if qu != nil && qu.Func != nil && qu.Func.Name == hybridFunc { + hasHybrid = true + break + } + } + if !hasHybrid { + return nil + } + + // Guard against the (extremely unlikely) case of a user variable colliding with + // the synthetic channel names we generate, which would otherwise produce a + // confusing "defined multiple times" error the user can't act on. Variables can + // be defined in nested blocks too, so check the whole query tree, not just roots. + for _, qu := range res.Query { + if v, ok := findReservedHybridVar(qu); ok { + return fmt.Errorf("variable %q uses the reserved prefix %q (used internally by hybrid)", + v, hybridVarPrefix) + } + } + + expanded := make([]*GraphQuery, 0, len(res.Query)+2) + hybridIdx := 0 + for _, qu := range res.Query { + if qu == nil || qu.Func == nil || qu.Func.Name != hybridFunc { + expanded = append(expanded, qu) + continue + } + blocks, err := expandHybridBlock(qu, hybridIdx) + if err != nil { + return err + } + hybridIdx++ + expanded = append(expanded, blocks...) + } + res.Query = expanded + return nil +} + +// validateFuseBlocks rejects fuse() blocks whose modifiers would be silently +// ignored. fuse() is computed coordinator-side from resolved value variables and is +// skipped by ProcessGraph, so an @filter, ordering, pagination, cascade, or child +// selection on the fuse block itself never executes. It must also be bound to a +// variable (that is the only way its result is consumed), and its channel list must +// not repeat a variable (a duplicate would double that channel's contribution). +func validateFuseBlocks(res *Result) error { + for _, qu := range res.Query { + if qu == nil || qu.Func == nil || qu.Func.Name != fuseFunc { + continue + } + if qu.Var == "" { + return fmt.Errorf("fuse must be assigned to a variable, e.g. " + + "`f as var(func: fuse(a, b))`; consume it with uid(f) / val(f)") + } + if qu.Filter != nil || len(qu.Order) > 0 || len(qu.Cascade) > 0 || + len(qu.Children) > 0 || qu.Args["first"] != "" || qu.Args["offset"] != "" { + return fmt.Errorf("fuse: @filter, ordering, pagination, cascade, and child " + + "blocks are not supported on a fuse() block; apply them on the " + + "consuming uid(var) block (use topk: to bound fused results)") + } + seen := make(map[string]bool, len(qu.Func.NeedsVar)) + for _, nv := range qu.Func.NeedsVar { + if seen[nv.Name] { + return fmt.Errorf("fuse: duplicate channel variable %q", nv.Name) + } + seen[nv.Name] = true + } + } + return nil +} + +// findReservedHybridVar walks a query block and its children for any variable using +// the reserved hybrid prefix, returning the first one found. +func findReservedHybridVar(qu *GraphQuery) (string, bool) { + if qu == nil { + return "", false + } + if strings.HasPrefix(qu.Var, hybridVarPrefix) { + return qu.Var, true + } + for _, ch := range qu.Children { + if v, ok := findReservedHybridVar(ch); ok { + return v, true + } + } + return "", false +} + +// expandHybridBlock turns a single hybrid() block into [bm25, similar_to, fuse]. +func expandHybridBlock(qu *GraphQuery, idx int) ([]*GraphQuery, error) { + if qu.Var == "" { + return nil, fmt.Errorf("hybrid must be assigned to a variable, e.g. " + + "`x as var(func: hybrid(textPred, \"query\", vecPred, $vec))`") + } + // The rewrite builds fresh channel/fuse blocks and discards everything else on the + // original hybrid block, so silently accepting modifiers here would drop them. + // Reject them explicitly: apply @filter/order/cascade on the outer uid(var) block. + if qu.Filter != nil || len(qu.Children) > 0 || len(qu.Order) > 0 || len(qu.Cascade) > 0 || + qu.Args["first"] != "" || qu.Args["offset"] != "" { + return nil, fmt.Errorf("hybrid: @filter, ordering, pagination, cascade, and child " + + "blocks are not supported on a hybrid() block; bind it to a variable and apply " + + "them on the uid(var) block (use topk: to bound the fused candidate set)") + } + fn := qu.Func + textPred := fn.Attr + if textPred == "" { + return nil, fmt.Errorf("hybrid: missing text predicate (first argument)") + } + + // hybrid has exactly three positional args in Args (the text predicate is the + // function Attr): queryText, vecPred and the query vector. Any further args are + // key/value option pairs appended by the parser. + const numPositional = 3 + if len(fn.Args) < numPositional { + return nil, fmt.Errorf("hybrid requires textPred, \"query text\", vecPred and a "+ + "query vector; got %d positional arguments", len(fn.Args)+1) + } + // Keep the whole Arg (not just its Value) so a parameterized text query keeps its + // IsDQLVar flag and gets substituted — just like the vector arg below. Taking only + // .Value would send the literal "$var" to bm25. + queryArg := fn.Args[0] + vecPred := fn.Args[1].Value + vecArg := fn.Args[2] + + // Options follow the positionals as key/value pairs; an odd remainder means a + // malformed option list rather than something to silently drop. + if (len(fn.Args)-numPositional)%2 != 0 { + return nil, fmt.Errorf("hybrid: malformed options (expected key:value pairs)") + } + + // Parse options: topk feeds similar_to's neighbor count; the rest feed fuse. + topk := hybridDefaultTopK + var fuseArgs []Arg + for i := numPositional; i+1 < len(fn.Args); i += 2 { + key := strings.ToLower(fn.Args[i].Value) + val := fn.Args[i+1] + if key == hybridTopKOption { + topk = val.Value + continue + } + fuseArgs = append(fuseArgs, Arg{Value: key}, val) + } + + // topk feeds similar_to's neighbor count and bm25's `first`; a non-positive value + // reaches HNSW as expectedNeighbors<=0 and panics the search. Reject it up front. + if n, err := strconv.Atoi(topk); err != nil || n <= 0 { + return nil, fmt.Errorf("hybrid: topk must be a positive integer, got %q", topk) + } + + chanBM25 := fmt.Sprintf("%s%d_bm25", hybridVarPrefix, idx) + chanVec := fmt.Sprintf("%s%d_vec", hybridVarPrefix, idx) + + // Bound the bm25 channel to the same topk candidate budget as the vector channel + // so a broad text query does not score and materialize the entire corpus before + // fusion. bm25 honors `first` with WAND top-k early termination. + bm25Block := &GraphQuery{ + Alias: "var", + Var: chanBM25, + Func: &Function{Name: "bm25", Attr: textPred, Args: []Arg{queryArg}}, + Args: map[string]string{"first": topk}, + } + simBlock := &GraphQuery{ + Alias: "var", + Var: chanVec, + Func: &Function{Name: similarToFn, Attr: vecPred, Args: []Arg{{Value: topk}, vecArg}}, + Args: map[string]string{}, + } + + channels := []VarContext{ + {Name: chanBM25, Typ: ValueVar}, + {Name: chanVec, Typ: ValueVar}, + } + fuseBlock := &GraphQuery{ + Alias: "var", + Var: qu.Var, + Func: &Function{Name: fuseFunc, NeedsVar: channels, Args: fuseArgs}, + NeedsVar: channels, + Args: map[string]string{}, + } + + return []*GraphQuery{bm25Block, simBlock, fuseBlock}, nil +} diff --git a/dql/parser.go b/dql/parser.go index 0dd6e1db7ac..32a63424cc1 100644 --- a/dql/parser.go +++ b/dql/parser.go @@ -29,8 +29,15 @@ const ( countFunc = "count" uidInFunc = "uid_in" similarToFn = "similar_to" + fuseFunc = "fuse" + hybridFunc = "hybrid" ) +// fuseOptionKeys is the set of named options accepted by the fuse() function. +var fuseOptionKeys = map[string]struct{}{ + "method": {}, "k": {}, "weights": {}, "normalize": {}, "topk": {}, +} + var ( errExpandType = "expand is only compatible with type filters" ) @@ -699,6 +706,18 @@ func ParseWithNeedVars(r Request, needVars []string) (res Result, rerr error) { } } + // Expand any hybrid() sugar blocks into explicit bm25 + similar_to channel + // blocks plus a fuse() block before variable collection and dependency checks. + if err := rewriteHybridBlocks(&res); err != nil { + return res, err + } + // fuse() executes coordinator-side and never reaches ProcessGraph, so block + // modifiers (@filter, ordering, pagination, children) would be silently ignored. + // Reject them, and reject duplicate channels, instead of degrading silently. + if err := validateFuseBlocks(&res); err != nil { + return res, err + } + if len(res.Query) != 0 { res.QueryVars = make([]*Vars, 0, len(res.Query)) for i := range res.Query { @@ -1701,7 +1720,8 @@ func validFuncName(name string) bool { switch name { case "regexp", "anyofterms", "allofterms", "alloftext", "anyoftext", "ngram", - "has", "uid", "uid_in", "anyof", "allof", "type", "match", "similar_to": + "has", "uid", "uid_in", "anyof", "allof", "type", "match", "similar_to", "bm25", + fuseFunc, hybridFunc: return true } return false @@ -1749,6 +1769,10 @@ L: if function.Name == similarToFn { similarToOptSeen = make(map[string]struct{}) } + var fuseOptSeen map[string]struct{} + if function.Name == fuseFunc || function.Name == hybridFunc { + fuseOptSeen = make(map[string]struct{}) + } if _, ok := tryParseItemType(it, itemLeftRound); !ok { return nil, it.Errorf("Expected ( after func name [%s]", function.Name) } @@ -1980,6 +2004,46 @@ L: // Disallow extra positional args after (k, vec). Options must be named. return nil, itemInFunc.Errorf("Expected named parameter in similar_to options (e.g. ef: 64)") } + + // fuse(v1, v2, ..., method: "rrf", k: 60) collects bare value-variable + // names as channels (handled by the fuseFunc case in the NeedsVar switch + // below) and key:value pairs as named options. An itemName followed by a + // colon is an option; otherwise it falls through to bare-name handling. + // hybrid() shares the same named options (its positional args are handled + // by the generic bare-name path). + if itemInFunc.Typ == itemName && + (function.Name == fuseFunc || function.Name == hybridFunc) { + next, ok := it.PeekOne() + if ok && next.Typ == itemColon { + key := strings.ToLower(collectName(it, itemInFunc.Val)) + if _, valid := fuseOptionKeys[key]; !valid { + return nil, itemInFunc.Errorf("Unknown option %q in fuse", key) + } + if _, exists := fuseOptSeen[key]; exists { + return nil, itemInFunc.Errorf("Duplicate key %q in fuse options", key) + } + fuseOptSeen[key] = struct{}{} + if ok := trySkipItemTyp(it, itemColon); !ok { + return nil, it.Errorf("Expected colon(:) after %s", key) + } + if !it.Next() { + return nil, it.Errorf("Expected value for %s", key) + } + valItem := it.Item() + if valItem.Typ != itemName { + return nil, valItem.Errorf("Expected value for %s", key) + } + v := strings.Trim(collectName(it, valItem.Val), " \t") + uq, err := unquoteIfQuoted(v) + if err != nil { + return nil, err + } + function.Args = append(function.Args, Arg{Value: key}, Arg{Value: uq}) + expectArg = false + continue + } + } + if itemInFunc.Typ != itemName { return nil, itemInFunc.Errorf("Expected arg after func [%s], but got item %v", function.Name, itemInFunc) @@ -2029,7 +2093,7 @@ L: // Unlike other functions, uid function has no attribute, everything is args. switch { case len(function.Attr) == 0 && function.Name != uidFunc && - function.Name != typFunc: + function.Name != typFunc && function.Name != fuseFunc: if strings.ContainsRune(itemInFunc.Val, '"') { return nil, itemInFunc.Errorf("Attribute in function"+ @@ -2047,7 +2111,7 @@ L: } function.Lang = val expectLang = false - case function.Name != uidFunc: + case function.Name != uidFunc && function.Name != fuseFunc: // For UID function. we set g.UID function.Args = append(function.Args, Arg{Value: val}) } @@ -2058,6 +2122,12 @@ L: expectArg = false switch function.Name { + case fuseFunc: + // fuse(v1, v2, ...) takes scored value variables as channels. + function.NeedsVar = append(function.NeedsVar, VarContext{ + Name: val, + Typ: ValueVar, + }) case valueFunc: // E.g. @filter(gt(val(a), 10)) function.NeedsVar = append(function.NeedsVar, VarContext{ @@ -2099,10 +2169,15 @@ L: } } - if function.Name != uidFunc && function.Name != typFunc && len(function.Attr) == 0 { + if function.Name != uidFunc && function.Name != typFunc && function.Name != fuseFunc && + len(function.Attr) == 0 { return nil, it.Errorf("Got empty attr for function: [%s]", function.Name) } + if function.Name == fuseFunc && len(function.NeedsVar) == 0 { + return nil, it.Errorf("fuse function requires at least one value variable channel") + } + if function.Name == typFunc && len(function.Args) != 1 { return nil, it.Errorf("type function only supports one argument. Got: %v", function.Args) } diff --git a/posting/bm25.go b/posting/bm25.go new file mode 100644 index 00000000000..a2e72ac7d74 --- /dev/null +++ b/posting/bm25.go @@ -0,0 +1,318 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package posting + +import ( + "context" + "encoding/binary" + "sync/atomic" + + ostats "go.opencensus.io/stats" + + "github.com/dgraph-io/dgraph/v25/protos/pb" + "github.com/dgraph-io/dgraph/v25/tok" + "github.com/dgraph-io/dgraph/v25/x" +) + +// BM25Posting is a single materialized entry of a BM25 term posting list: the +// document UID together with the term frequency and document length decoded from +// the posting value. +type BM25Posting struct { + Uid uint64 + TF uint32 + DocLen uint32 +} + +// ReadBM25TermPostings materializes the postings of a BM25 term's standard index +// list at readTs into a UID-ascending slice, decoding (tf, docLen) from each +// posting value. getList reads a posting list for a key (e.g. LocalCache.Get), +// keeping the value encoding encapsulated in this package. +func ReadBM25TermPostings(getList func(key []byte) (*List, error), attr, encodedTerm string, + readTs uint64) ([]BM25Posting, error) { + key := x.BM25IndexKey(attr, encodedTerm) + pl, err := getList(key) + if err != nil { + return nil, err + } + var out []BM25Posting + err = pl.Iterate(readTs, 0, func(p *pb.Posting) error { + tf, docLen, ok := decodeBM25Value(p.Value) + if !ok { + // Corrupt/truncated posting value: skip rather than inject a bogus + // zero-frequency match that would silently distort scoring. + return nil + } + out = append(out, BM25Posting{Uid: p.Uid, TF: tf, DocLen: docLen}) + return nil + }) + if err != nil { + return nil, err + } + return out, nil +} + +// NumBM25StatsBuckets is the number of buckets the BM25 corpus statistics (document +// count and total term count) are sharded across, keyed by uid%NumBM25StatsBuckets. +// Sharding spreads the read-modify-write contention of stats maintenance across +// independent posting lists so that concurrent mutations on different documents +// rarely conflict, while same-bucket updates still conflict (and retry) — avoiding +// lost updates. A single hot stats key would serialize all writes to the predicate. +// Exported so the bulk loader buckets corpus statistics identically to the live and +// rebuild paths. +const NumBM25StatsBuckets = 32 + +// numBM25StatsBuckets is the unexported alias retained for readability within this +// package. +const numBM25StatsBuckets = NumBM25StatsBuckets + +// EncodeBM25Value packs a posting's term frequency and document length the same way the +// live index path does, for the bulk loader to write BM25 term postings in the standard +// format. See encodeBM25Value. +func EncodeBM25Value(tf, docLen uint32) []byte { return encodeBM25Value(tf, docLen) } + +// EncodeBM25Stats encodes corpus statistics (document count, total term count) for the +// bulk loader to write the per-bucket stats postings in the standard format. See +// encodeBM25Stats. +func EncodeBM25Stats(docCount, totalTerms uint64) []byte { + return encodeBM25Stats(docCount, totalTerms) +} + +// DecodeBM25Stats decodes a stats posting value written by EncodeBM25Stats. The bulk +// reducer uses it to fold per-document stats postings into bucket totals. +func DecodeBM25Stats(b []byte) (docCount, totalTerms uint64) { + return decodeBM25Stats(b) +} + +// encodeBM25Value packs a posting's term frequency and document length into the +// posting Value as two unsigned varints. Storing the document length alongside the +// term frequency makes scoring read (tf, docLen) in a single posting access — no +// separate document-length list (which would be a write-hot key) and no random +// per-candidate lookup at query time. The document length is duplicated across a +// document's unique terms, but a document's postings are always rewritten together +// on update, so they stay consistent. +func encodeBM25Value(tf, docLen uint32) []byte { + buf := make([]byte, binary.MaxVarintLen32*2) + n := binary.PutUvarint(buf, uint64(tf)) + n += binary.PutUvarint(buf[n:], uint64(docLen)) + return buf[:n] +} + +// decodeBM25Value reverses encodeBM25Value. ok is false when the input does not +// hold two complete varints (e.g. a truncated or corrupt posting value), so the +// caller can skip it rather than silently scoring it as a zero-frequency match. +func decodeBM25Value(b []byte) (tf, docLen uint32, ok bool) { + tf64, n := binary.Uvarint(b) + if n <= 0 { + return 0, 0, false + } + docLen64, m := binary.Uvarint(b[n:]) + if m <= 0 { + return 0, 0, false + } + return uint32(tf64), uint32(docLen64), true +} + +// encodeBM25Stats encodes corpus statistics (document count, total term count) as +// two unsigned varints. +func encodeBM25Stats(docCount, totalTerms uint64) []byte { + buf := make([]byte, binary.MaxVarintLen64*2) + n := binary.PutUvarint(buf, docCount) + n += binary.PutUvarint(buf[n:], totalTerms) + return buf[:n] +} + +// decodeBM25Stats reverses encodeBM25Stats. It returns (0, 0) on malformed input. +func decodeBM25Stats(b []byte) (docCount, totalTerms uint64) { + docCount, n := binary.Uvarint(b) + if n <= 0 { + return 0, 0 + } + totalTerms, m := binary.Uvarint(b[n:]) + if m <= 0 { + return docCount, 0 + } + return docCount, totalTerms +} + +// addBM25TermPosting writes (op=SET) or removes (op=DEL) the posting for the given +// (term, uid) pair in the term's standard index posting list. On SET the posting's +// Value packs (tf, docLen); on DEL only the UID matters. The posting is a REF +// posting (ValueId set) that carries a Value — List.encode retains such postings +// through rollup (see the len(p.Value) > 0 clause there). +func (txn *Txn) addBM25TermPosting(ctx context.Context, attr, term string, uid uint64, + tf, docLen uint32, op pb.DirectedEdge_Op) error { + encodedTerm := string([]byte{tok.IdentBM25}) + term + key := x.BM25IndexKey(attr, encodedTerm) + plist, err := txn.cache.GetFromDelta(key) + if err != nil { + return err + } + edge := &pb.DirectedEdge{ + ValueId: uid, + Attr: attr, + Op: op, + } + if op != pb.DirectedEdge_DEL { + edge.Value = encodeBM25Value(tf, docLen) + edge.ValueType = pb.Posting_BINARY + } + if err := plist.addMutation(ctx, txn, edge); err != nil { + return err + } + ostats.Record(ctx, x.NumEdges.M(1)) + return nil +} + +// bm25StatsAccum is a concurrency-safe per-bucket accumulator of corpus-statistics +// deltas. Index rebuild routes its per-document stats updates here (across many +// goroutines) instead of the read-modify-write counter, then flushes the buckets once +// as a single writer (see flush). This avoids the undercount that the streaming +// rebuild's independent per-thread caches and periodic resets would otherwise cause, +// where last-write-wins on the value posting drops every thread's partial total but +// one. +type bm25StatsAccum struct { + count [numBM25StatsBuckets]atomic.Int64 + terms [numBM25StatsBuckets]atomic.Int64 +} + +func newBM25StatsAccum() *bm25StatsAccum { return &bm25StatsAccum{} } + +// add records a document's contribution in its bucket (uid%numBM25StatsBuckets). +func (a *bm25StatsAccum) add(uid uint64, docCountDelta, totalTermsDelta int64) { + bucket := uid % numBM25StatsBuckets + a.count[bucket].Add(docCountDelta) + a.terms[bucket].Add(totalTermsDelta) +} + +// flush writes the accumulated absolute totals into the stats posting lists for attr +// through txn, one SET value posting per non-empty bucket. It is a single-writer +// operation (one txn writing all buckets), so there is no lost-update window; the +// caller commits txn. Buckets are written as absolute SETs because a rebuild deletes +// the prior stats first, so the buckets start empty. +func (a *bm25StatsAccum) flush(ctx context.Context, txn *Txn, attr string) error { + for bucket := 0; bucket < numBM25StatsBuckets; bucket++ { + docCount := a.count[bucket].Load() + totalTerms := a.terms[bucket].Load() + if docCount <= 0 && totalTerms <= 0 { + continue + } + key := x.BM25StatsKey(attr, bucket) + plist, err := txn.cache.GetFromDelta(key) + if err != nil { + return err + } + edge := &pb.DirectedEdge{ + Attr: attr, + Value: encodeBM25Stats(uint64(docCount), uint64(totalTerms)), + ValueType: pb.Posting_BINARY, + Op: pb.DirectedEdge_SET, + } + if err := plist.addMutation(ctx, txn, edge); err != nil { + return err + } + } + return nil +} + +// updateBM25Stats applies (docCountDelta, totalTermsDelta) to the bucketed corpus +// statistics for attr. The bucket is selected by uid%numBM25StatsBuckets. The +// running totals are stored as a single value posting per bucket; the read at +// txn.StartTs sees this transaction's own earlier writes (read-your-own-writes), +// so multiple documents in the same transaction that land in the same bucket +// accumulate correctly. +func (txn *Txn) updateBM25Stats(ctx context.Context, attr string, uid uint64, + docCountDelta, totalTermsDelta int64) error { + // During index rebuild, accumulate into the shared accumulator rather than the + // read-modify-write counter (see Txn.bm25Acc). The rebuild flushes the buckets + // once at the end as a single writer. + if txn.bm25Acc != nil { + txn.bm25Acc.add(uid, docCountDelta, totalTermsDelta) + return nil + } + // Serialize the whole read-modify-write against sibling goroutines applying + // other edges of the SAME transaction (see Txn.bm25StatsMu): their entities can + // share this uid%32 bucket, and concurrent RMWs would drop deltas. + txn.bm25StatsMu.Lock() + defer txn.bm25StatsMu.Unlock() + + bucket := int(uid % numBM25StatsBuckets) + key := x.BM25StatsKey(attr, bucket) + // Stats are maintained by read-modify-write: we must read the committed total + // from disk (and merge this transaction's own writes), not just the in-memory + // delta. GetFromDelta skips disk and is only safe for write-only index mutations, + // so each transaction would otherwise overwrite the bucket instead of + // accumulating across transactions. Get reads committed state. + plist, err := txn.cache.Get(key) + if err != nil { + return err + } + + var docCount, totalTerms uint64 + val, err := plist.Value(txn.StartTs) + switch err { + case nil: + if data, ok := val.Value.([]byte); ok { + docCount, totalTerms = decodeBM25Stats(data) + } + case ErrNoValue: + // No stats yet for this bucket; start from zero. + default: + return err + } + + docCount = applyBM25Delta(docCount, docCountDelta) + totalTerms = applyBM25Delta(totalTerms, totalTermsDelta) + + edge := &pb.DirectedEdge{ + Attr: attr, + Value: encodeBM25Stats(docCount, totalTerms), + ValueType: pb.Posting_BINARY, + Op: pb.DirectedEdge_SET, + } + return plist.addMutation(ctx, txn, edge) +} + +// applyBM25Delta adds a signed delta to an unsigned counter, clamping at zero. +func applyBM25Delta(v uint64, delta int64) uint64 { + if delta >= 0 { + return v + uint64(delta) + } + dec := uint64(-delta) + if dec > v { + return 0 + } + return v - dec +} + +// ReadBM25Stats sums the bucketed corpus statistics for attr at readTs, returning +// the document count and total term count. avgDL = totalTerms / docCount. The +// getList closure reads a posting list for a key (e.g. LocalCache.Get) so the +// caller controls caching and the read timestamp. +func ReadBM25Stats(getList func(key []byte) (*List, error), attr string, + readTs uint64) (docCount, totalTerms uint64, err error) { + for b := 0; b < numBM25StatsBuckets; b++ { + key := x.BM25StatsKey(attr, b) + pl, perr := getList(key) + if perr != nil { + return 0, 0, perr + } + val, verr := pl.Value(readTs) + if verr == ErrNoValue { + continue + } + if verr != nil { + return 0, 0, verr + } + data, ok := val.Value.([]byte) + if !ok || len(data) == 0 { + continue + } + dc, tt := decodeBM25Stats(data) + docCount += dc + totalTerms += tt + } + return docCount, totalTerms, nil +} diff --git a/posting/bm25_bench_test.go b/posting/bm25_bench_test.go new file mode 100644 index 00000000000..a417f55906c --- /dev/null +++ b/posting/bm25_bench_test.go @@ -0,0 +1,213 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package posting + +import ( + "context" + "math/rand" + "sync" + "testing" + + "github.com/dgraph-io/dgraph/v25/x" +) + +// Benchmarks for the fixed per-query BM25 overheads (campaign Q35). +// +// Every bm25() query unconditionally calls ReadBM25Stats (worker/task.go, +// handleBM25Search), which performs NumBM25StatsBuckets (32) posting-list +// reads regardless of query selectivity — a df=10 query pays the same floor +// as a df=1M one. The ReadBM25Stats benchmarks measure that floor against the +// package's badger-backed pstore three ways: straight store reads (no cache), +// a fresh LocalCache per call (what a real query does), and a warm LocalCache +// (repeat reads within one query). +// +// The codec round-trip benches pin the cost of the varint encodings that sit +// on every posting materialization (tf, docLen) and every stats read/write +// (docCount, totalTerms); decode of the posting value is the per-posting hot +// path inside ReadBM25TermPostings. +// +// The companion Q34 benchmark (scored HNSW search overhead) is deliberately +// deferred: it requires a populated HNSW index, which is not available at +// this package's unit level. + +var bm25BenchOnce sync.Once + +// bm25BenchStatsSetup commits one transaction that writes deterministic corpus +// statistics into every one of the 32 stats buckets for the benchmark +// predicate, mirroring the commit path of TestBM25StatsAccumulateAcrossTxns. +// It returns the namespaced attribute and a timestamp at which the committed +// stats are visible. +func bm25BenchStatsSetup(tb testing.TB) (attr string, readTs uint64) { + attr = x.AttrInRootNamespace("bm25benchstats") + const startTs, commitTs = 50001, 50002 + readTs = commitTs + 1 + bm25BenchOnce.Do(func() { + ctx := context.Background() + rng := rand.New(rand.NewSource(42)) // deterministic corpus + txn := Oracle().RegisterStartTs(startTs) + txn.cache = NewLocalCache(startTs) + // 320 documents with uids 1..320: every bucket (uid%32) receives + // exactly ten documents, so all 32 bucket keys hold a value posting. + for uid := uint64(1); uid <= 320; uid++ { + docLen := int64(rng.Intn(100) + 1) + if err := txn.updateBM25Stats(ctx, attr, uid, 1, docLen); err != nil { + tb.Fatal(err) + } + } + txn.Update() + txn.UpdateCachedKeys(commitTs) + writer := NewTxnWriter(pstore) + if err := txn.CommitToDisk(writer, commitTs); err != nil { + tb.Fatal(err) + } + if err := writer.Flush(); err != nil { + tb.Fatal(err) + } + }) + return attr, readTs +} + +// checkBM25BenchStats guards against the benchmark timing an accidental no-op: +// the corpus is 320 docs, and totalTerms must be non-zero. +func checkBM25BenchStats(b *testing.B, docCount, totalTerms uint64) { + b.Helper() + if docCount != 320 || totalTerms == 0 { + b.Fatalf("unexpected stats: docCount=%d totalTerms=%d", docCount, totalTerms) + } +} + +// BenchmarkReadBM25StatsStore measures the uncached stats floor: each op reads +// all 32 bucket posting lists straight from the badger-backed store. +func BenchmarkReadBM25StatsStore(b *testing.B) { + attr, readTs := bm25BenchStatsSetup(b) + get := func(k []byte) (*List, error) { return GetNoStore(k, readTs) } + var dc, tt uint64 + var err error + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + dc, tt, err = ReadBM25Stats(get, attr, readTs) + if err != nil { + b.Fatal(err) + } + } + b.StopTimer() + checkBM25BenchStats(b, dc, tt) + b.ReportMetric(numBM25StatsBuckets, "buckets_read/op") +} + +// BenchmarkReadBM25StatsFreshLocalCache is the production-faithful shape: a +// query starts with an empty LocalCache, so all 32 bucket reads miss and fall +// through to the store, plus the cache bookkeeping. +func BenchmarkReadBM25StatsFreshLocalCache(b *testing.B) { + attr, readTs := bm25BenchStatsSetup(b) + var dc, tt uint64 + var err error + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + lc := NewLocalCache(readTs) + dc, tt, err = ReadBM25Stats(lc.Get, attr, readTs) + if err != nil { + b.Fatal(err) + } + } + b.StopTimer() + checkBM25BenchStats(b, dc, tt) + b.ReportMetric(numBM25StatsBuckets, "buckets_read/op") +} + +// BenchmarkReadBM25StatsWarmLocalCache measures the repeat-read floor when all +// 32 bucket lists are already materialized in the query's LocalCache. +func BenchmarkReadBM25StatsWarmLocalCache(b *testing.B) { + attr, readTs := bm25BenchStatsSetup(b) + lc := NewLocalCache(readTs) + // Prime the cache so the timed loop measures only warm reads. + if _, _, err := ReadBM25Stats(lc.Get, attr, readTs); err != nil { + b.Fatal(err) + } + var dc, tt uint64 + var err error + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + dc, tt, err = ReadBM25Stats(lc.Get, attr, readTs) + if err != nil { + b.Fatal(err) + } + } + b.StopTimer() + checkBM25BenchStats(b, dc, tt) + b.ReportMetric(numBM25StatsBuckets, "buckets_read/op") +} + +// BenchmarkBM25ValueRoundTrip measures the (tf, docLen) posting-value codec: +// one encode plus one decode per op, over varint widths from 1 to 5 bytes. +func BenchmarkBM25ValueRoundTrip(b *testing.B) { + var sink uint32 + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + tf := uint32(i%127 + 1) // 1-byte varint + dl := uint32(i%(1<<20)) + 128 // 2..3-byte varint + buf := encodeBM25Value(tf, dl) + gtf, gdl, ok := decodeBM25Value(buf) + if !ok || gtf != tf || gdl != dl { + b.Fatalf("round trip mismatch: (%d,%d) -> (%d,%d,%v)", tf, dl, gtf, gdl, ok) + } + sink += gtf + gdl + } + if sink == 0 { + b.Fatal("sink is zero; codec produced no data") + } +} + +// BenchmarkBM25ValueDecode isolates the decode half — the per-posting hot path +// of ReadBM25TermPostings, executed once per materialized posting. +func BenchmarkBM25ValueDecode(b *testing.B) { + // Deterministic pre-encoded values covering the varint width range. + rng := rand.New(rand.NewSource(42)) + const n = 1024 + encoded := make([][]byte, n) + for i := range encoded { + encoded[i] = encodeBM25Value(uint32(rng.Intn(1<<16)), uint32(rng.Intn(1<<24))) + } + var sink uint32 + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + tf, dl, ok := decodeBM25Value(encoded[i%n]) + if !ok { + b.Fatal("decode failed on valid input") + } + sink += tf + dl + } + if sink == 0 { + b.Fatal("sink is zero; decode produced no data") + } + b.ReportMetric(1, "postings_decoded/op") +} + +// BenchmarkBM25StatsRoundTrip measures the (docCount, totalTerms) stats codec: +// one encode plus one decode per op — the payload of every bucket read/write. +func BenchmarkBM25StatsRoundTrip(b *testing.B) { + var sink uint64 + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + dc := uint64(i%1_000_000 + 1) + tt := dc * 37 // plausible avgDL ~37 + buf := encodeBM25Stats(dc, tt) + gdc, gtt := decodeBM25Stats(buf) + if gdc != dc || gtt != tt { + b.Fatalf("round trip mismatch: (%d,%d) -> (%d,%d)", dc, tt, gdc, gtt) + } + sink += gdc + gtt + } + if sink == 0 { + b.Fatal("sink is zero; codec produced no data") + } +} diff --git a/posting/bm25_test.go b/posting/bm25_test.go new file mode 100644 index 00000000000..9e4913f555c --- /dev/null +++ b/posting/bm25_test.go @@ -0,0 +1,283 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package posting + +import ( + "bytes" + "context" + "math" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + "github.com/dgraph-io/dgraph/v25/protos/pb" + "github.com/dgraph-io/dgraph/v25/x" +) + +// TestBM25DropClearsStats guards against orphaned corpus statistics when the BM25 +// index is dropped or rebuilt. BM25 term postings live under the IdentBM25 token +// prefix, but the stats buckets live under a separate reserved token prefix. The +// drop/rebuild machinery deletes by token prefix, so unless the stats prefix is +// also returned by prefixesToDeleteTokensFor, the stats survive a drop and then +// double-count when the index is rebuilt on top of them. +func TestBM25DropClearsStats(t *testing.T) { + attr := x.AttrInRootNamespace("bm25dropstats") + prefixes, err := prefixesToDeleteTokensFor(attr, "bm25", false) + require.NoError(t, err) + + // Every stats bucket key must be covered by one of the deletion prefixes. + for bucket := 0; bucket < 3; bucket++ { + statsKey := x.BM25StatsKey(attr, bucket) + covered := false + for _, p := range prefixes { + if bytes.HasPrefix(statsKey, p) { + covered = true + break + } + } + require.Truef(t, covered, + "stats bucket %d key not covered by any bm25 deletion prefix (orphaned on drop)", bucket) + } +} + +func TestBM25ValueCodecRoundTrip(t *testing.T) { + cases := [][2]uint32{{0, 0}, {1, 1}, {3, 12}, {7, 200}, {65535, 1 << 20}, {1 << 24, 1 << 24}} + for _, c := range cases { + tf, dl, ok := decodeBM25Value(encodeBM25Value(c[0], c[1])) + require.True(t, ok) + require.Equal(t, c[0], tf) + require.Equal(t, c[1], dl) + } + // Malformed/truncated input is reported as invalid so callers can skip it. + _, _, ok := decodeBM25Value(nil) + require.False(t, ok) + _, _, ok = decodeBM25Value([]byte{0x80}) // varint continuation byte with no terminator + require.False(t, ok) +} + +// TestBM25ValueSurvivesRollup verifies the linchpin of the BM25 redesign: a REF +// index posting that carries a packed (tf, docLen) value is retained — value and +// all — through rollup, instead of being collapsed to a UID-only Pack entry (the +// default behavior for REF postings before the len(p.Value) > 0 retention clause +// in List.encode). +func TestBM25ValueSurvivesRollup(t *testing.T) { + attr := x.AttrInRootNamespace("bm25rollup") + encodedTerm := string([]byte{0x10}) + "fox" // IdentBM25 || term + key := x.BM25IndexKey(attr, encodedTerm) + + docs := []struct { + uid uint64 + tf uint32 + docLen uint32 + }{ + {uid: 5, tf: 3, docLen: 12}, + {uid: 9, tf: 1, docLen: 40}, + {uid: 100, tf: 7, docLen: 200}, + } + + ts := uint64(1) + for _, d := range docs { + l, err := GetNoStore(key, ts) + require.NoError(t, err) + edge := &pb.DirectedEdge{ + ValueId: d.uid, + Attr: attr, + Value: encodeBM25Value(d.tf, d.docLen), + ValueType: pb.Posting_BINARY, + } + addMutation(t, l, edge, Set, ts, ts+1, false) + ts += 2 + } + + // Force a rollup and decode the resulting posting list directly. + l, err := getNew(key, pstore, math.MaxUint64, false) + require.NoError(t, err) + kvs, err := l.Rollup(nil, math.MaxUint64) + require.NoError(t, err) + require.NotEmpty(t, kvs) + + var plist pb.PostingList + require.NoError(t, proto.Unmarshal(kvs[0].Value, &plist)) + + got := make(map[uint64][2]uint32) + for _, p := range plist.Postings { + tf, docLen, ok := decodeBM25Value(p.Value) + require.True(t, ok) + got[p.Uid] = [2]uint32{tf, docLen} + } + for _, d := range docs { + v, ok := got[d.uid] + require.Truef(t, ok, "uid %d posting missing after rollup (value stripped?)", d.uid) + require.Equal(t, d.tf, v[0], "tf for uid %d", d.uid) + require.Equal(t, d.docLen, v[1], "docLen for uid %d", d.uid) + } + + // Reading the list back materializes the same (uid, tf, docLen) triples. + posts, err := ReadBM25TermPostings(func(k []byte) (*List, error) { + return getNew(k, pstore, math.MaxUint64, false) + }, attr, encodedTerm, math.MaxUint64) + require.NoError(t, err) + require.Len(t, posts, len(docs)) + for _, p := range posts { + require.Equal(t, got[p.Uid][0], p.TF) + require.Equal(t, got[p.Uid][1], p.DocLen) + } +} + +// TestBM25StatsConflictKeyValueIndependent guards the invariant that makes the +// corpus-stats read-modify-write counter safe under concurrent live transactions: +// two overlapping transactions that read the same bucket base and write DIFFERENT +// resulting totals must still be detected as conflicting (so one retries instead of +// silently losing an update). This holds because on a scalar (non-list) predicate +// fingerprintEdge returns MaxUint64 for every untagged value, so all stats writes to +// a bucket share the conflict key getKey(statsKey, MaxUint64) — independent of the +// value bytes. It is precisely why @index(bm25) is rejected on list predicates, where +// fingerprintEdge would become value-dependent and let differing totals slip past +// conflict detection. +func TestBM25StatsConflictKeyValueIndependent(t *testing.T) { + attr := x.AttrInRootNamespace("bm25statsconflict") + key := x.BM25StatsKey(attr, 3) + pk, err := x.Parse(key) + require.NoError(t, err) + + // Mirror addMutationInternal: a value posting (no ValueId) gets its ValueId set + // from fingerprintEdge before the conflict key is computed. + mkEdge := func(val []byte) *pb.DirectedEdge { + e := &pb.DirectedEdge{Attr: attr, Value: val, ValueType: pb.Posting_BINARY, Op: pb.DirectedEdge_SET} + if NewPosting(e).PostingType != pb.Posting_REF { + e.ValueId = fingerprintEdge(e) + } + return e + } + e1 := mkEdge(encodeBM25Stats(10, 100)) + e2 := mkEdge(encodeBM25Stats(11, 137)) // different totals + require.Equal(t, uint64(math.MaxUint64), e1.ValueId, + "scalar untagged stats values must carry the MaxUint64 sentinel ValueId") + require.Equal(t, e1.ValueId, e2.ValueId, "stats ValueId must not depend on the value bytes") + + ck1 := GetConflictKey(pk, key, e1) + ck2 := GetConflictKey(pk, key, e2) + require.NotZero(t, ck1, "stats writes must register a conflict key") + require.Equal(t, ck1, ck2, + "two writers to the same stats bucket must share a conflict key regardless of the value") +} + +// TestBM25StatsRebuildAccumulator covers the fix for stats undercounting during index +// rebuild. The streaming rebuild processes documents across many goroutines, each with +// its own transaction/cache that is periodically reset — so a per-transaction +// read-modify-write counter loses updates (the last writer's partial total wins on +// merge). Routing rebuild stats through a shared accumulator and flushing the buckets +// once (single writer) must reproduce the exact corpus totals. This test models the +// rebuild by feeding documents (several sharing a bucket) through SEPARATE +// transactions that all share one accumulator, then flushing and reading back. +func TestBM25StatsRebuildAccumulator(t *testing.T) { + ctx := context.Background() + attr := x.AttrInRootNamespace("bm25rebuildacc") + acc := newBM25StatsAccum() + + docs := []struct { + uid uint64 + dl int64 + }{{1, 10}, {2, 20}, {33, 5}, {65, 7}, {3, 8}, {35, 4}, {97, 9}, {4, 6}} + var wantCount, wantTerms int64 + for _, d := range docs { + txn := NewTxn(900) + txn.cache = NewLocalCache(900) + txn.bm25Acc = acc + require.NoError(t, txn.updateBM25Stats(ctx, attr, d.uid, 1, d.dl)) + wantCount++ + wantTerms += d.dl + } + + // Single-writer finalize: flush the accumulator through one transaction and commit. + txn := Oracle().RegisterStartTs(901) + txn.cache = NewLocalCache(901) + require.NoError(t, acc.flush(ctx, txn, attr)) + txn.Update() + txn.UpdateCachedKeys(902) + writer := NewTxnWriter(pstore) + require.NoError(t, txn.CommitToDisk(writer, 902)) + require.NoError(t, writer.Flush()) + + get := func(k []byte) (*List, error) { return GetNoStore(k, 903) } + dc, tt, err := ReadBM25Stats(get, attr, 903) + require.NoError(t, err) + require.Equal(t, uint64(wantCount), dc, "rebuilt doc count must equal the total across all transactions") + require.Equal(t, uint64(wantTerms), tt, "rebuilt total terms must equal the total across all transactions") +} + +// TestBM25StatsBucketed verifies that bucketed corpus statistics accumulate +// correctly across documents (including two documents that hash to the same +// bucket, exercising in-transaction read-your-own-writes) and that deletes +// subtract correctly. +func TestBM25StatsBucketed(t *testing.T) { + ctx := context.Background() + attr := x.AttrInRootNamespace("bm25stats") + ts := uint64(101) + txn := Oracle().RegisterStartTs(ts) + + // uid 1 and uid 33 both fall in bucket 1 (mod 32), exercising same-bucket + // accumulation within a single transaction. + docs := []struct { + uid uint64 + dl int64 + }{{1, 10}, {2, 20}, {33, 5}, {64, 7}, {100, 8}} + + var wantCount, wantTerms int64 + for _, d := range docs { + require.NoError(t, txn.updateBM25Stats(ctx, attr, d.uid, 1, d.dl)) + wantCount++ + wantTerms += d.dl + } + + get := func(k []byte) (*List, error) { return txn.cache.GetFromDelta(k) } + dc, tt, err := ReadBM25Stats(get, attr, ts) + require.NoError(t, err) + require.Equal(t, uint64(wantCount), dc) + require.Equal(t, uint64(wantTerms), tt) + + // Delete uid 2: docCount and totalTerms drop accordingly. + require.NoError(t, txn.updateBM25Stats(ctx, attr, 2, -1, -20)) + dc, tt, err = ReadBM25Stats(get, attr, ts) + require.NoError(t, err) + require.Equal(t, uint64(wantCount-1), dc) + require.Equal(t, uint64(wantTerms-20), tt) +} + +// TestBM25StatsAccumulateAcrossTxns verifies that stats accumulate across +// separately-committed transactions (not just within one). This guards against +// the read-modify-write reading only the in-memory delta instead of committed +// disk state, which would make each transaction overwrite its bucket and collapse +// the corpus document count. +func TestBM25StatsAccumulateAcrossTxns(t *testing.T) { + ctx := context.Background() + attr := x.AttrInRootNamespace("bm25statsxtxn") + + // Two documents in the SAME bucket (uid 5 and uid 37 → bucket 5), committed in + // two separate transactions. + commitDoc := func(startTs, commitTs, uid uint64, docLen int64) { + txn := Oracle().RegisterStartTs(startTs) + txn.cache = NewLocalCache(startTs) + require.NoError(t, txn.updateBM25Stats(ctx, attr, uid, 1, docLen)) + txn.Update() + txn.UpdateCachedKeys(commitTs) + writer := NewTxnWriter(pstore) + require.NoError(t, txn.CommitToDisk(writer, commitTs)) + require.NoError(t, writer.Flush()) + } + + commitDoc(201, 202, 5, 10) + commitDoc(203, 204, 37, 6) + + // A fresh reader at a later ts must see BOTH documents (count 2, terms 16), + // not just the most recently committed one. + get := func(k []byte) (*List, error) { return GetNoStore(k, 205) } + dc, tt, err := ReadBM25Stats(get, attr, 205) + require.NoError(t, err) + require.Equal(t, uint64(2), dc, "doc count must accumulate across transactions") + require.Equal(t, uint64(16), tt, "total terms must accumulate across transactions") +} diff --git a/posting/index.go b/posting/index.go index ae6c3352a44..6412a59960c 100644 --- a/posting/index.go +++ b/posting/index.go @@ -68,6 +68,10 @@ func indexTokens(ctx context.Context, info *indexMutationInfo) ([]string, error) var tokens []string for _, it := range info.tokenizers { + // BM25 tokenizer is handled separately in addBM25IndexMutations. + if it.Identifier() == tok.IdentBM25 { + continue + } toks, err := tok.BuildTokens(sv.Value, tok.GetTokenizerForLang(it, lang)) if err != nil { return tokens, err @@ -179,6 +183,17 @@ func (txn *Txn) addIndexMutations(ctx context.Context, info *indexMutationInfo) } } + // Check if any tokenizer is BM25 and handle separately. + for _, it := range info.tokenizers { + if _, ok := tok.GetTokenizerForLang(it, info.edge.GetLang()).(tok.BM25Tokenizer); ok { + if err := txn.addBM25IndexMutations(ctx, info); err != nil { + return []*pb.DirectedEdge{}, err + } + // Continue to process remaining non-BM25 tokenizers below. + continue + } + } + tokens, err := indexTokens(ctx, info) if err != nil { // This data is not indexable @@ -215,6 +230,58 @@ func (txn *Txn) addIndexMutation(ctx context.Context, edge *pb.DirectedEdge, tok return nil } +// addBM25IndexMutations handles index mutations for the BM25 tokenizer. Unlike +// other tokenizers, each BM25 index posting carries a value that packs the term +// frequency together with the document length (see encodeBM25Value). The postings +// are written through the standard delta path (plist.addMutation), so BM25 rides +// Dgraph's normal posting-list machinery — MVCC, deltas, rollup, splits, backup — +// with no separate storage path. Corpus statistics (document count and total term +// count, from which the average document length is derived) are kept in bucketed +// stats posting lists keyed by uid%numBM25StatsBuckets to avoid a single write-hot +// key while preserving conflict detection per bucket. +// +// Updates are driven entirely by the caller (AddMutationWithIndex), which issues a +// DEL for the previous value followed by a SET for the new one. The DEL re-tokenizes +// the old value and removes its postings and stats contribution; the SET adds the new +// ones. We therefore never need to detect updates here. +func (txn *Txn) addBM25IndexMutations(ctx context.Context, info *indexMutationInfo) error { + attr := info.edge.Attr + uid := info.edge.Entity + lang := info.edge.GetLang() + + schemaType, err := schema.State().TypeOf(attr) + if err != nil || !schemaType.IsScalar() { + return errors.Errorf("Cannot BM25 index attribute %s of type object.", attr) + } + + sv, err := types.Convert(info.val, schemaType) + if err != nil { + return err + } + + bm25Tok := tok.BM25Tokenizer{} + termFreqs, docLen, err := bm25Tok.TokensWithFrequency(sv.Value, lang) + if err != nil { + return err + } + + // Skip documents that tokenize to zero terms (e.g., all stopwords). + if docLen == 0 { + return nil + } + + for term, tf := range termFreqs { + if err := txn.addBM25TermPosting(ctx, attr, term, uid, tf, docLen, info.op); err != nil { + return err + } + } + + if info.op == pb.DirectedEdge_DEL { + return txn.updateBM25Stats(ctx, attr, uid, -1, -int64(docLen)) + } + return txn.updateBM25Stats(ctx, attr, uid, 1, int64(docLen)) +} + // countParams is sent to updateCount function. It is used to update the count index. // It deletes the uid from the key corresponding to and adds it // to . @@ -666,6 +733,14 @@ func prefixesToDeleteTokensFor(attr, tokenizerName string, hasLang bool) ([][]by prefix = append(prefix, tokenizer.Identifier()) prefixes = append(prefixes, prefix) + // BM25 stores corpus statistics under a reserved token prefix separate from its + // term-posting (IdentBM25) prefix, so deleting only the tokenizer-identifier + // prefix above would orphan the stats and double-count on rebuild. Delete the + // stats prefix (and its split variant) alongside the term postings. + if tokenizerName == "bm25" { + prefixes = append(prefixes, x.BM25StatsPrefix(attr, false), x.BM25StatsPrefix(attr, true)) + } + return prefixes, nil } @@ -678,6 +753,20 @@ type rebuilder struct { // The posting list passed here is the on disk version. It is not coming // from the LRU cache. fn func(uid uint64, pl *List, txn *Txn) ([]*pb.DirectedEdge, error) + + // bm25Acc, when non-nil, collects BM25 corpus statistics across the rebuild's + // per-thread transactions so they can be flushed once as a single writer. Set by + // rebuildTokIndex when a BM25 tokenizer is being rebuilt. See Txn.bm25Acc. + bm25Acc *bm25StatsAccum +} + +// newRebuildTxn creates a transaction for the streaming rebuild, propagating the +// shared BM25 stats accumulator (nil for non-BM25 rebuilds) so per-thread stats +// updates are collected rather than lost across cache resets. +func (r *rebuilder) newRebuildTxn() *Txn { + txn := NewTxn(r.startTs) + txn.bm25Acc = r.bm25Acc + return txn } func (r *rebuilder) RunWithoutTemp(ctx context.Context) error { @@ -941,7 +1030,7 @@ func (r *rebuilder) Run(ctx context.Context) error { txns := make([]*Txn, maxThreadIds) for i := range txns { - txns[i] = NewTxn(r.startTs) + txns[i] = r.newRebuildTxn() } stream.FinishThread = func(threadId int) (*bpb.KVList, error) { @@ -961,7 +1050,7 @@ func (r *rebuilder) Run(ctx context.Context) error { } kvs = append(kvs, &kv) } - txns[threadId] = NewTxn(r.startTs) + txns[threadId] = r.newRebuildTxn() return &bpb.KVList{Kv: kvs}, nil } @@ -1020,7 +1109,7 @@ func (r *rebuilder) Run(ctx context.Context) error { kvs = append(kvs, &kv) } - txns[threadId] = NewTxn(r.startTs) + txns[threadId] = r.newRebuildTxn() return &bpb.KVList{Kv: kvs}, nil } @@ -1041,6 +1130,25 @@ func (r *rebuilder) Run(ctx context.Context) error { if err := stream.Orchestrate(ctx); err != nil { return err } + // Flush the BM25 corpus statistics accumulated across all rebuild threads as a + // single writer. They are written into the temp store as ordinary delta postings + // so the second phase rolls them up into pstore at r.startTs alongside the term + // postings — no separate commit path, and no per-thread last-write-wins loss. + if r.bm25Acc != nil { + flushTxn := NewTxn(r.startTs) + flushTxn.cache = NewLocalCache(r.startTs) + if err := r.bm25Acc.flush(ctx, flushTxn, r.attr); err != nil { + return err + } + flushTxn.Update() + for key, data := range flushTxn.cache.deltas { + counter++ + e := &badger.Entry{Key: []byte(key), Value: data, UserMeta: BitDeltaPosting} + if err := tmpWriter.SetEntryAt(e, counter); err != nil { + return errors.Wrap(err, "error writing bm25 stats to temp index") + } + } + } if err := tmpWriter.Flush(); err != nil { return err } @@ -1446,6 +1554,14 @@ func rebuildTokIndex(ctx context.Context, rb *IndexRebuild) error { pk := x.ParsedKey{Attr: rb.Attr} builder := rebuilder{attr: rb.Attr, prefix: pk.DataPrefix(), startTs: rb.StartTs} + // BM25 corpus statistics must be aggregated across the rebuild's per-thread + // transactions and flushed once; a per-thread read-modify-write counter would + // undercount. Route stats through a shared accumulator when rebuilding bm25. + for _, tokenizer := range tokenizers { + if tokenizer.Identifier() == tok.IdentBM25 { + builder.bm25Acc = newBM25StatsAccum() + } + } builder.fn = func(uid uint64, pl *List, txn *Txn) ([]*pb.DirectedEdge, error) { edge := pb.DirectedEdge{Attr: rb.Attr, Entity: uid} edges := []*pb.DirectedEdge{} diff --git a/posting/list.go b/posting/list.go index 1c0c7a0fc55..610eaf5b5c2 100644 --- a/posting/list.go +++ b/posting/list.go @@ -1627,7 +1627,14 @@ func (l *List) encode(out *rollupOutput, readTs uint64, split bool) error { } enc.Add(p.Uid) - if p.Facets != nil || p.PostingType != pb.Posting_REF { + // Retain the full posting (not just its UID in the Pack) whenever it + // carries facets, is not a plain UID reference, or carries a value. + // BM25 index postings are REF postings that pack (term-frequency, + // doc-length) into Value; without the len(p.Value) > 0 clause that + // value would be stripped at rollup, silently losing all term + // frequencies. This mirrors how faceted postings already coexist in + // both Pack (UID) and Postings (payload). + if p.Facets != nil || p.PostingType != pb.Posting_REF || len(p.Value) > 0 { plist.Postings = append(plist.Postings, p) } return nil diff --git a/posting/mvcc.go b/posting/mvcc.go index 81c5e375553..108cdfc3b3e 100644 --- a/posting/mvcc.go +++ b/posting/mvcc.go @@ -318,6 +318,7 @@ func (txn *Txn) CommitToDisk(writer *TxnWriter, commitTs uint64) error { return err } } + return nil } diff --git a/posting/oracle.go b/posting/oracle.go index d7c3837b4b2..3b718b05784 100644 --- a/posting/oracle.go +++ b/posting/oracle.go @@ -54,6 +54,22 @@ type Txn struct { lastUpdate time.Time cache *LocalCache // This pointer does not get modified. + + // bm25Acc, when non-nil, redirects BM25 corpus-statistics updates into a shared + // accumulator instead of the per-transaction read-modify-write counter. Index + // rebuild sets this on its per-thread transactions so stats survive the streaming + // rebuild's independent caches and periodic resets, which would otherwise drop + // updates and undercount the corpus. nil on normal live transactions. + bm25Acc *bm25StatsAccum + + // bm25StatsMu serializes the BM25 corpus-statistics read-modify-write within + // this transaction. A large proposal is applied by parallel goroutines batched + // so that no two batches share an (attr, entity) DATA key — but stats buckets + // are keyed by uid%numBM25StatsBuckets, so DIFFERENT entities in different + // batches legitimately hit the SAME bucket concurrently, and an unserialized + // RMW loses one goroutine's delta (last write wins). Cross-transaction safety + // is unaffected (value-independent conflict keys still serialize commits). + bm25StatsMu sync.Mutex } // struct to implement Txn interface from vector-indexer diff --git a/query/common_test.go b/query/common_test.go index e36211f7a18..64a88bd865e 100644 --- a/query/common_test.go +++ b/query/common_test.go @@ -390,6 +390,15 @@ func populateCluster(dc dgraphapi.Cluster) { testSchema += "\ndescription: string @index(ngram) ." } + // BM25 indexing - uses same version gate as ngram for now + if ngramSupport { + testSchema += "\ndescription_bm25: string @index(bm25) ." + // A parallel vector predicate on the same documents enables hybrid-search + // (fuse) tests that combine BM25 with vector similarity. Gated together with + // BM25 so both channels are always present for those tests. + testSchema += "\ndescription_vec: float32vector @index(hnsw(metric:\"euclidean\")) ." + } + setSchema(testSchema) err = addTriplesToCluster(` @@ -1007,4 +1016,31 @@ func populateCluster(dc dgraphapi.Cluster) { <415> "Linguistic analysis helps understand text meaning" . `) x.Panic(err) + + // Add data for BM25 tests - uses separate predicate to avoid conflicts + err = addTriplesToCluster(` + <501> "The quick brown fox jumps over the lazy dog" . + <502> "A quick brown fox leaps over a sleeping dog" . + <503> "fox fox fox" . + <504> "The lazy dog sleeps under the warm sun all day long in the garden" . + <505> "Dogs are loyal companions to humans and families everywhere" . + <506> "Quick movements help foxes catch their prey in the wild" . + <507> "Brown foxes are quick and agile animals in the forest" . + `) + x.Panic(err) + + // Vector embeddings on the same documents (dims = [fox, dog, quick, brown]). + // Chosen so a "pure fox" query vector [3,0,0,0] ranks 503 first, then the + // fox/quick docs (506,507,501,502), and the dog-only docs (504,505) last — + // letting hybrid (fuse) tests observe both channels and union semantics. + err = addTriplesToCluster(` + <501> "[1.0, 1.0, 1.0, 1.0]" . + <502> "[1.0, 1.0, 1.0, 1.0]" . + <503> "[3.0, 0.0, 0.0, 0.0]" . + <504> "[0.0, 2.0, 0.0, 0.0]" . + <505> "[0.0, 2.0, 0.0, 0.0]" . + <506> "[1.0, 0.0, 1.0, 0.0]" . + <507> "[1.0, 0.0, 1.0, 1.0]" . + `) + x.Panic(err) } diff --git a/query/fuse.go b/query/fuse.go new file mode 100644 index 00000000000..f82bd48e918 --- /dev/null +++ b/query/fuse.go @@ -0,0 +1,450 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package query + +import ( + "container/heap" + "math" + "sort" + "strconv" + "strings" + + "github.com/dgraph-io/dgraph/v25/dql" + "github.com/dgraph-io/dgraph/v25/protos/pb" + "github.com/dgraph-io/dgraph/v25/types" + "github.com/pkg/errors" +) + +// This file implements the pure, I/O-free core of native hybrid search: combining +// several already-scored result sets ("channels") into a single ranked list. +// +// A channel is a uid->score map produced by an upstream ranker bound to a DQL value +// variable (e.g. bm25(...) or similar_to(...)). Fusion is a coordinator-side +// operation over resolved value variables; everything here is deterministic and +// independent of storage, sharding, and the query pipeline so it can be tested in +// isolation. The query-layer adapter (see populateUidValVar) converts value +// variables into fuseChannels and the result back into a value variable. + +// fusionMethod selects how channel scores are combined. +type fusionMethod int + +const ( + // fusionRRF is Reciprocal Rank Fusion: each channel contributes 1/(k+rank), + // where rank is the 1-based position of the uid within that channel. Robust to + // heterogeneous score scales because it uses ranks, not raw scores. + fusionRRF fusionMethod = iota + // fusionLinear is a weighted sum of (optionally normalized) raw scores. + fusionLinear +) + +// linearNormalize selects score normalization for fusionLinear. +type linearNormalize int + +const ( + // normalizeMax divides each channel's scores by that channel's maximum absolute + // score and clamps the result to [0,1]. Clamping matters under the union's + // missing-uid=0 convention: a signed metric (cosine/dot) can retrieve a document + // with a negative similarity, and without the clamp that document would fuse below + // a document the channel never retrieved at all (which contributes 0). Clamping + // negative similarities to 0 makes "retrieved but dissimilar" tie with "not + // retrieved" instead of ranking beneath it. BM25 and euclidean (1/(1+d)) scores are + // already >=0, so they are unaffected. + normalizeMax linearNormalize = iota + // normalizeNone uses raw scores as-is (the caller asserts comparability). + normalizeNone +) + +// defaultRRFK is the conventional RRF rank constant. Larger k flattens the +// contribution of top ranks; 60 is the widely used default. +const defaultRRFK = 60.0 + +// fuseChannel is one scored input to fusion. scores maps uid -> raw score with the +// convention that higher is always better (all Dgraph rankers surface +// higher-is-better scores). weight applies only to fusionLinear. +type fuseChannel struct { + scores map[uint64]float64 + weight float64 +} + +// fuseOpts configures a fusion run. +type fuseOpts struct { + method fusionMethod + k float64 // RRF rank constant; <=0 falls back to defaultRRFK. + normalize linearNormalize // linear only. + topk int // if >0, truncate output to the top topk results. +} + +// scoredUid is a uid paired with its fused score. +type scoredUid struct { + uid uint64 + score float64 +} + +// fuseChannels combines channels into a single ranked list, sorted by fused score +// descending and tie-broken by uid ascending. The candidate set is the UNION of all +// channels' uids (outer join): a uid missing from a channel simply receives no +// contribution from it (RRF: as if rank = infinity; linear: as if score = 0). It is +// never dropped and never produces NaN. +func fuseChannels(channels []fuseChannel, opts fuseOpts) []scoredUid { + var fused map[uint64]float64 + switch opts.method { + case fusionLinear: + fused = fuseLinear(channels, opts.normalize) + default: + fused = fuseRRF(channels, opts.k) + } + + if opts.topk > 0 && len(fused) > opts.topk { + return topKFused(fused, opts.topk) + } + out := make([]scoredUid, 0, len(fused)) + for uid, s := range fused { + out = append(out, scoredUid{uid: uid, score: s}) + } + sortFusedDesc(out) + return out +} + +// sortFusedDesc sorts by score descending, ties broken by uid ascending. +func sortFusedDesc(out []scoredUid) { + sort.Slice(out, func(i, j int) bool { + if out[i].score != out[j].score { + return out[i].score > out[j].score + } + return out[i].uid < out[j].uid + }) +} + +// fusedHeap is a worst-at-root heap of scoredUid under (score desc, uid asc) +// ranking: the root is the entry that would be evicted first. +type fusedHeap []scoredUid + +func (h fusedHeap) Len() int { return len(h) } +func (h fusedHeap) Less(i, j int) bool { + // Root must be the WORST entry: lower score first; among equal scores the + // higher uid is worse (uid asc wins ties in the final ranking). + if h[i].score != h[j].score { + return h[i].score < h[j].score + } + return h[i].uid > h[j].uid +} +func (h fusedHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *fusedHeap) Push(x interface{}) { *h = append(*h, x.(scoredUid)) } +func (h *fusedHeap) Pop() interface{} { + old := *h + n := len(old) + x := old[n-1] + *h = old[:n-1] + return x +} + +// topKFused selects the topk entries by (score desc, uid asc) without sorting the +// whole fused map: a size-bounded worst-at-root heap makes selection +// O(n log topk) instead of O(n log n), and the n-sized output slice is never +// allocated. +func topKFused(fused map[uint64]float64, topk int) []scoredUid { + // Fill phase appends directly (heap.Push would box each scoredUid through + // interface{}, one allocation per element); a single Init establishes the + // heap invariant once the buffer is full. + h := make(fusedHeap, 0, topk) + for uid, s := range fused { + cand := scoredUid{uid: uid, score: s} + if len(h) < topk { + h = append(h, cand) + if len(h) == topk { + heap.Init(&h) + } + continue + } + // Replace the root if the candidate outranks the current worst. + if root := h[0]; cand.score > root.score || + (cand.score == root.score && cand.uid < root.uid) { + h[0] = cand + heap.Fix(&h, 0) + } + } + out := []scoredUid(h) + sortFusedDesc(out) + return out +} + +// channelOrder returns the channel's uids sorted by score descending, tie-broken +// by uid ascending — the 1-based rank of a uid is its index+1. The tie-break +// matches the deterministic ordering used elsewhere in the codebase (bm25/HNSW +// sorted()), so equal-scored uids rank stably by uid. Scores are copied out of +// the map once so the sort comparator does slice reads, not map lookups (two map +// probes per comparison dominated the sort's cost at 100k+ uids). +func channelOrder(c fuseChannel) []uint64 { + pairs := make([]scoredUid, 0, len(c.scores)) + for uid, s := range c.scores { + pairs = append(pairs, scoredUid{uid: uid, score: s}) + } + sortFusedDesc(pairs) + uids := make([]uint64, len(pairs)) + for i, p := range pairs { + uids[i] = p.uid + } + return uids +} + +// fuseRRF computes (weighted) Reciprocal Rank Fusion over the channels. Each +// channel contributes weight * 1/(k+rank); with the default weight of 1.0 this is +// standard RRF, and per-channel weights let callers bias channels under either +// fusion method rather than silently ignoring weights for rrf. Ranks come from +// iterating the sorted order positionally — the intermediate uid->rank map was a +// measured hotspot (~70% of 8-channel RRF wall time at 100k uids/channel). +func fuseRRF(channels []fuseChannel, k float64) map[uint64]float64 { + if k <= 0 || math.IsNaN(k) || math.IsInf(k, 0) { + k = defaultRRFK + } + fused := make(map[uint64]float64) + for _, c := range channels { + for i, uid := range channelOrder(c) { + fused[uid] += c.weight * (1.0 / (k + float64(i+1))) + } + } + return fused +} + +// channelMaxAbs returns the maximum finite absolute score in a channel, used to +// max-normalize heterogeneous score scales. Non-finite scores are ignored; +// returns 0 for an empty channel (callers treat a 0 denominator as "contribute 0"). +func channelMaxAbs(c fuseChannel) float64 { + var maxAbs float64 + for _, s := range c.scores { + if math.IsNaN(s) || math.IsInf(s, 0) { + continue + } + if a := math.Abs(s); a > maxAbs { + maxAbs = a + } + } + return maxAbs +} + +// fuseLinear computes a weighted sum of (optionally max-normalized) raw scores. A +// uid missing from a channel contributes 0 from that channel. A channel whose +// maximum absolute score is 0 (all zeros / empty) contributes 0 rather than +// dividing by zero. +func fuseLinear(channels []fuseChannel, normalize linearNormalize) map[uint64]float64 { + denoms := make([]float64, len(channels)) + for i, c := range channels { + if normalize == normalizeMax { + denoms[i] = channelMaxAbs(c) + } else { + denoms[i] = 1.0 + } + } + + fused := make(map[uint64]float64) + for i, c := range channels { + denom := denoms[i] + for uid, s := range c.scores { + norm := s + if normalize == normalizeMax { + if denom == 0 { + norm = 0 + } else { + norm = s / denom + } + // Clamp negatives to the missing-uid baseline (0) so a retrieved but + // dissimilar document never fuses below one the channel never retrieved. + if norm < 0 { + norm = 0 + } + } + fused[uid] += c.weight * norm + } + } + // Ensure uids that appear only in zero-contribution channels are still present + // in the union (e.g. an all-zero max-normalized channel). The loop above already + // inserts them with a running sum (possibly 0), so the union is complete. + return fused +} + +// --- Query-layer adapter ----------------------------------------------------- +// +// The functions below bridge the pure fusion core to DQL value variables. They +// parse the fuse() options, read each channel's scores from the already-resolved +// variable map, run fusion, and return a varValue carrying both the union uid set +// and the fused uid->score map (the same shape the bm25 ranker binds). + +// parseFuseOpts extracts fuse() options from the function's key/value arg pairs and +// returns the resolved fuseOpts plus the optional per-channel linear weights +// (nil when unspecified). numChannels is used to validate the weights count. +func parseFuseOpts(args []dql.Arg, numChannels int) (fuseOpts, []float64, error) { + opts := fuseOpts{method: fusionRRF, k: defaultRRFK, normalize: normalizeMax} + var weights []float64 + + for i := 0; i+1 < len(args); i += 2 { + key := strings.ToLower(args[i].Value) + val := args[i+1].Value + switch key { + case "method": + switch strings.ToLower(val) { + case "rrf": + opts.method = fusionRRF + case "linear": + opts.method = fusionLinear + default: + return opts, nil, errors.Errorf("fuse: unknown method %q (want rrf or linear)", val) + } + case "k": + k, err := strconv.ParseFloat(val, 64) + if err != nil || k <= 0 || math.IsNaN(k) || math.IsInf(k, 0) { + return opts, nil, errors.Errorf("fuse: k must be a positive finite number, got %q", val) + } + opts.k = k + case "normalize": + switch strings.ToLower(val) { + case "max": + opts.normalize = normalizeMax + case "none": + opts.normalize = normalizeNone + default: + return opts, nil, errors.Errorf("fuse: unknown normalize %q (want max or none)", val) + } + case "topk": + tk, err := strconv.Atoi(val) + if err != nil || tk < 0 { + return opts, nil, errors.Errorf("fuse: topk must be a non-negative integer, got %q", val) + } + opts.topk = tk + case "weights": + parts := strings.Split(val, ",") + weights = make([]float64, 0, len(parts)) + for _, p := range parts { + w, err := strconv.ParseFloat(strings.TrimSpace(p), 64) + if err != nil || math.IsNaN(w) || math.IsInf(w, 0) { + return opts, nil, errors.Errorf("fuse: invalid weight %q", p) + } + weights = append(weights, w) + } + default: + return opts, nil, errors.Errorf("fuse: unknown option %q", key) + } + } + + if weights != nil && len(weights) != numChannels { + return opts, nil, errors.Errorf("fuse: weights count (%d) must match channel count (%d)", + len(weights), numChannels) + } + return opts, weights, nil +} + +// computeFuse reads the channel value variables named in the fuse function's +// NeedsVar from doneVars, runs fusion, and returns a varValue with the union uid +// set and the fused uid->score map. +func computeFuse(args []dql.Arg, needsVar []dql.VarContext, + doneVars map[string]varValue, sgPath []*SubGraph) (varValue, error) { + + if len(needsVar) == 0 { + return varValue{}, errors.Errorf("fuse: requires at least one value variable channel") + } + + opts, weights, err := parseFuseOpts(args, len(needsVar)) + if err != nil { + return varValue{}, err + } + + channels := make([]fuseChannel, len(needsVar)) + for i, nv := range needsVar { + v, ok := doneVars[nv.Name] + switch { + case !ok: + // The dependency scheduler guarantees every channel block has run and + // populated doneVars before this fuse block. A genuinely absent channel + // therefore signals an internal invariant violation rather than an empty + // result — surface it instead of silently degrading the fusion. + return varValue{}, errors.Errorf("fuse: channel %q was not produced", nv.Name) + // Note: ShardedMap.IsEmpty() only detects nil (a fresh NewShardedMap has 30 + // empty shards), so entry-count emptiness must use Len() == 0. + case v.Vals == nil || v.Vals.Len() == 0: + // A uid variable carries matched uids but no scores. Treating it as an + // empty channel would silently drop its uids from the fused ranking, so + // reject it explicitly — fusion needs scored channels. + if v.Uids != nil && len(v.Uids.GetUids()) > 0 { + return varValue{}, errors.Errorf("fuse: channel %q is a uid variable "+ + "without scores; use a ranker such as bm25 or similar_to", nv.Name) + } + // A channel that ran but matched nothing is a valid empty channel: it + // contributes nothing but must not drop the other channels' results. + channels[i] = fuseChannel{scores: map[uint64]float64{}, weight: 1.0} + default: + scores, err := scoresFromVar(v, nv.Name) + if err != nil { + return varValue{}, err + } + channels[i] = fuseChannel{scores: scores, weight: 1.0} + } + if weights != nil { + channels[i].weight = weights[i] + } + } + + fused := fuseChannels(channels, opts) + + out := varValue{Vals: types.NewShardedMap(), path: sgPath} + uids := make([]uint64, len(fused)) + for i, r := range fused { + uids[i] = r.uid + out.Vals.Set(r.uid, types.Val{Tid: types.FloatID, Value: r.score}) + } + // Emit the uid set in ascending order — the Dgraph value-variable contract (the + // same one bm25 follows): Uids is the unordered candidate set used by uid(var), + // and the fused score lives in Vals. Callers recover ranked order with + // `orderdesc: val(var)`. When `topk` is set, fuseChannels has already selected + // the top-k by fused score before this ascending sort, so top-k + orderdesc is + // correct. (Sorting here by score would break uid(var) set semantics.) + sort.Slice(uids, func(i, j int) bool { return uids[i] < uids[j] }) + out.Uids = &pb.List{Uids: uids} + return out, nil +} + +// scoresFromVar extracts a uid->float64 score map from a value variable. The +// variable must carry numeric scores (as bm25/similar_to bind); a uid variable +// without scores cannot be a fusion channel. +func scoresFromVar(v varValue, name string) (map[uint64]float64, error) { + scores := make(map[uint64]float64, v.Vals.Len()) + var convErr error + err := v.Vals.Iterate(func(uid uint64, val types.Val) error { + var s float64 + // bm25 and similar_to bind FloatID scores directly; take that fast path so a + // non-finite value (which types.Convert rejects) is dropped rather than + // mis-reported as "non-numeric". Other numeric types go through Convert. + if val.Tid == types.FloatID { + f, ok := val.Value.(float64) + if !ok { + convErr = errors.Errorf("fuse: channel %q has a malformed score value", name) + return convErr + } + s = f + } else { + f, err := types.Convert(val, types.FloatID) + if err != nil { + convErr = errors.Errorf("fuse: channel %q has non-numeric scores; use a "+ + "ranker such as bm25 or similar_to", name) + return convErr + } + s = f.Value.(float64) + } + // Drop non-finite scores so they can never poison fusion: a NaN/Inf would + // break the sort comparator's strict-weak-ordering and propagate through + // linear sums. A uid dropped here simply doesn't participate via this channel. + if math.IsNaN(s) || math.IsInf(s, 0) { + return nil + } + scores[uid] = s + return nil + }) + if convErr != nil { + return nil, convErr + } + if err != nil { + return nil, err + } + return scores, nil +} diff --git a/query/fuse_bench_test.go b/query/fuse_bench_test.go new file mode 100644 index 00000000000..8d7ecb50f3b --- /dev/null +++ b/query/fuse_bench_test.go @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package query + +import ( + "fmt" + "math/rand" + "testing" +) + +// Q32 fusion-overhead benchmarks (pure unit level, no cluster). +// +// fuseChannels makes three to four full passes over the channel union with no +// early bound: channelRanks sorts every channel (RRF), fuseChannels materializes +// and sorts the full union, and topk truncation happens only AFTER the full +// sort. These benchmarks quantify that cost across channel count, channel size, +// overlap, fusion method, and topk, plus an isolation benchmark for +// channelRanks (the rank-map allocation is the suspected hotspot). +// +// Corpora are deterministic: every channel set is generated from a constant +// rand seed, so numbers are comparable across runs. + +const fuseBenchSeed = 42 + +// buildFuseBenchChannels builds numChannels channels of `size` uids each, where +// a fraction `overlap` of every channel's uids is drawn from a pool shared by +// ALL channels and the remainder is unique to that channel. Scores are uniform +// random in (0, 10). The union cardinality is therefore +// shared + numChannels*(size-shared). +func buildFuseBenchChannels(numChannels, size int, overlap float64) []fuseChannel { + rng := rand.New(rand.NewSource(fuseBenchSeed)) + shared := int(float64(size) * overlap) + + // Shared pool: uids 1..shared (uid 0 is invalid in Dgraph). + sharedUids := make([]uint64, shared) + for i := range sharedUids { + sharedUids[i] = uint64(i + 1) + } + + channels := make([]fuseChannel, numChannels) + for ci := 0; ci < numChannels; ci++ { + scores := make(map[uint64]float64, size) + for _, uid := range sharedUids { + scores[uid] = rng.Float64() * 10 + } + // Unique uids live in a per-channel range far above the shared pool. + base := uint64(ci+1) * 1_000_000_000 + for i := 0; i < size-shared; i++ { + scores[base+uint64(i)] = rng.Float64() * 10 + } + channels[ci] = fuseChannel{scores: scores, weight: 1.0} + } + return channels +} + +// fuseBenchUnionSize computes the union cardinality of the channels, counted in +// the harness because fuseChannels exposes no counter. +func fuseBenchUnionSize(channels []fuseChannel) int { + union := make(map[uint64]struct{}) + for _, c := range channels { + for uid := range c.scores { + union[uid] = struct{}{} + } + } + return len(union) +} + +// BenchmarkFuseChannels is the Q32 matrix: 2/4/8 channels x 1k/10k/100k uids +// per channel x 10%/90% overlap x rrf vs linear(normalizeMax) x topk 10 vs 0 +// (unbounded). union_uids/op is the fused candidate-set cardinality; +// fused_out/op is the emitted result length (post-topk). +func BenchmarkFuseChannels(b *testing.B) { + methods := []struct { + name string + opts fuseOpts + }{ + {"rrf", fuseOpts{method: fusionRRF, k: defaultRRFK}}, + {"linmax", fuseOpts{method: fusionLinear, normalize: normalizeMax}}, + } + + for _, numChannels := range []int{2, 4, 8} { + for _, size := range []int{1000, 10000, 100000} { + for _, overlap := range []float64{0.10, 0.90} { + channels := buildFuseBenchChannels(numChannels, size, overlap) + union := fuseBenchUnionSize(channels) + for _, m := range methods { + for _, topk := range []int{10, 0} { + opts := m.opts + opts.topk = topk + name := fmt.Sprintf("ch=%d/sz=%d/ov=%d/%s/topk=%d", + numChannels, size, int(overlap*100), m.name, topk) + b.Run(name, func(b *testing.B) { + b.ReportAllocs() + var res []scoredUid + for i := 0; i < b.N; i++ { + res = fuseChannels(channels, opts) + } + b.ReportMetric(float64(union), "union_uids/op") + b.ReportMetric(float64(len(res)), "fused_out/op") + }) + } + } + } + } + } +} + +// BenchmarkChannelOrder isolates the per-channel rank computation used by RRF: +// one full sort of the channel's uids plus a fresh map[uint64]int allocation +// per call (the suspected allocation hotspot inside fuseRRF's channel loop). +func BenchmarkChannelOrder(b *testing.B) { + for _, size := range []int{1000, 10000, 100000} { + channels := buildFuseBenchChannels(1, size, 0) + c := channels[0] + b.Run(fmt.Sprintf("sz=%d", size), func(b *testing.B) { + b.ReportAllocs() + var order []uint64 + for i := 0; i < b.N; i++ { + order = channelOrder(c) + } + b.ReportMetric(float64(len(order)), "ranks/op") + }) + } +} diff --git a/query/fuse_test.go b/query/fuse_test.go new file mode 100644 index 00000000000..b088c0c8202 --- /dev/null +++ b/query/fuse_test.go @@ -0,0 +1,407 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package query + +import ( + "math" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/dql" + "github.com/dgraph-io/dgraph/v25/types" +) + +// ch is a small helper to build a fusion channel from a uid->score map with a +// default weight of 1.0. +func ch(scores map[uint64]float64) fuseChannel { + return fuseChannel{scores: scores, weight: 1.0} +} + +// asMap collapses a fused result slice into a uid->score map for assertions that +// don't care about ordering. +func asMap(res []scoredUid) map[uint64]float64 { + m := make(map[uint64]float64, len(res)) + for _, r := range res { + m[r.uid] = r.score + } + return m +} + +func TestFuseRRF_BasicRanks(t *testing.T) { + // Channel A order: 10, 20, 30 (ranks 1,2,3) + // Channel B order: 30, 10, 40 (ranks 1,2,3) + a := ch(map[uint64]float64{10: 9.0, 20: 5.0, 30: 1.0}) + b := ch(map[uint64]float64{30: 0.9, 10: 0.5, 40: 0.1}) + + res := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionRRF, k: 60}) + got := asMap(res) + + const k = 60.0 + // uid 10: rank1 in A, rank2 in B + require.InDelta(t, 1/(k+1)+1/(k+2), got[10], 1e-9) + // uid 20: rank2 in A only + require.InDelta(t, 1/(k+2), got[20], 1e-9) + // uid 30: rank3 in A, rank1 in B + require.InDelta(t, 1/(k+3)+1/(k+1), got[30], 1e-9) + // uid 40: rank3 in B only + require.InDelta(t, 1/(k+3), got[40], 1e-9) +} + +func TestFuseRRF_OrderingAndUnion(t *testing.T) { + a := ch(map[uint64]float64{10: 9.0, 20: 5.0, 30: 1.0}) + b := ch(map[uint64]float64{30: 0.9, 10: 0.5, 40: 0.1}) + + res := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionRRF, k: 60}) + + // Union of all uids is present (outer join, not intersection). + require.Len(t, res, 4) + // Sorted by fused score descending. uid 10 and 30 both appear in both channels + // near the top; 10 is rank1+rank2, 30 is rank3+rank1 -> 10 slightly higher. + require.Equal(t, uint64(10), res[0].uid) + require.Equal(t, uint64(30), res[1].uid) + // Scores must be monotonically non-increasing. + for i := 1; i < len(res); i++ { + require.LessOrEqual(t, res[i].score, res[i-1].score) + } +} + +func TestFuseRRF_DefaultK(t *testing.T) { + a := ch(map[uint64]float64{1: 1.0}) + // k<=0 should fall back to the default of 60. + res := fuseChannels([]fuseChannel{a}, fuseOpts{method: fusionRRF, k: 0}) + require.InDelta(t, 1/(60.0+1), res[0].score, 1e-9) +} + +func TestFuseRRF_TieBreakByUidAscending(t *testing.T) { + // Equal scores within a channel -> lower uid gets the better (smaller) rank. + a := ch(map[uint64]float64{2: 5.0, 1: 5.0, 3: 5.0}) + res := fuseChannels([]fuseChannel{a}, fuseOpts{method: fusionRRF, k: 60}) + got := asMap(res) + // uid 1 rank1, uid 2 rank2, uid 3 rank3. + require.InDelta(t, 1/(60.0+1), got[1], 1e-9) + require.InDelta(t, 1/(60.0+2), got[2], 1e-9) + require.InDelta(t, 1/(60.0+3), got[3], 1e-9) + // Final output tie-broken by uid ascending when fused scores are equal. + require.Equal(t, uint64(1), res[0].uid) +} + +func TestFuseRRF_DisjointChannels(t *testing.T) { + a := ch(map[uint64]float64{1: 9.0, 2: 8.0}) + b := ch(map[uint64]float64{3: 9.0, 4: 8.0}) + res := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionRRF, k: 60}) + require.Len(t, res, 4) + got := asMap(res) + // Each uid scored only by its single channel rank. + require.InDelta(t, 1/(60.0+1), got[1], 1e-9) + require.InDelta(t, 1/(60.0+1), got[3], 1e-9) + require.InDelta(t, 1/(60.0+2), got[2], 1e-9) + require.InDelta(t, 1/(60.0+2), got[4], 1e-9) +} + +func TestFuseRRF_AppliesWeights(t *testing.T) { + // Weights must affect RRF (not only linear): a uid ranked #1 in a 2x-weighted + // channel should beat a uid ranked #1 in a unit-weighted channel. + heavy := fuseChannel{scores: map[uint64]float64{1: 9.0}, weight: 2.0} + light := fuseChannel{scores: map[uint64]float64{2: 9.0}, weight: 1.0} + res := fuseChannels([]fuseChannel{heavy, light}, fuseOpts{method: fusionRRF, k: 60}) + got := asMap(res) + require.InDelta(t, 2.0*(1/(60.0+1)), got[1], 1e-9) + require.InDelta(t, 1.0*(1/(60.0+1)), got[2], 1e-9) + require.Equal(t, uint64(1), res[0].uid, "heavier-weighted channel's top doc wins") +} + +func TestFuseRRF_DefaultWeightIsStandardRRF(t *testing.T) { + // With the default weight of 1.0, weighted RRF reduces to standard RRF. + a := ch(map[uint64]float64{10: 9.0, 20: 5.0}) + res := fuseChannels([]fuseChannel{a}, fuseOpts{method: fusionRRF, k: 60}) + got := asMap(res) + require.InDelta(t, 1/(60.0+1), got[10], 1e-9) + require.InDelta(t, 1/(60.0+2), got[20], 1e-9) +} + +func TestFuseLinear_MaxNormalizeAndWeights(t *testing.T) { + // BM25-ish scale vs cosine-ish scale. + text := fuseChannel{scores: map[uint64]float64{1: 10.0, 2: 5.0}, weight: 0.3} + vec := fuseChannel{scores: map[uint64]float64{1: 0.8, 2: 0.4, 3: 0.2}, weight: 0.7} + + res := fuseChannels([]fuseChannel{text, vec}, + fuseOpts{method: fusionLinear, normalize: normalizeMax}) + got := asMap(res) + + // max-normalize: text/10, vec/0.8. + // uid1: 0.3*(10/10) + 0.7*(0.8/0.8) = 0.3 + 0.7 = 1.0 + require.InDelta(t, 1.0, got[1], 1e-9) + // uid2: 0.3*(5/10) + 0.7*(0.4/0.8) = 0.15 + 0.35 = 0.5 + require.InDelta(t, 0.5, got[2], 1e-9) + // uid3: only vec: 0.7*(0.2/0.8) = 0.175 (text contributes 0, not NaN) + require.InDelta(t, 0.175, got[3], 1e-9) + + require.Equal(t, uint64(1), res[0].uid) +} + +func TestFuseLinear_NoNormalize(t *testing.T) { + a := fuseChannel{scores: map[uint64]float64{1: 2.0, 2: 1.0}, weight: 1.0} + b := fuseChannel{scores: map[uint64]float64{1: 3.0}, weight: 2.0} + res := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionLinear, normalize: normalizeNone}) + got := asMap(res) + // uid1: 1*2 + 2*3 = 8 ; uid2: 1*1 = 1 + require.InDelta(t, 8.0, got[1], 1e-9) + require.InDelta(t, 1.0, got[2], 1e-9) +} + +func TestFuseLinear_ZeroMaxChannelContributesZero(t *testing.T) { + // A channel whose scores are all zero must not divide-by-zero / NaN. + a := fuseChannel{scores: map[uint64]float64{1: 0.0, 2: 0.0}, weight: 1.0} + b := fuseChannel{scores: map[uint64]float64{1: 4.0}, weight: 1.0} + res := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionLinear, normalize: normalizeMax}) + got := asMap(res) + require.False(t, math.IsNaN(got[1])) + require.False(t, math.IsNaN(got[2])) + // uid1: a contributes 0, b contributes 4/4=1 -> 1.0 + require.InDelta(t, 1.0, got[1], 1e-9) + // uid2: only in a (all-zero) -> 0.0 + require.InDelta(t, 0.0, got[2], 1e-9) +} + +func TestFuseLinear_NegativeSimilarityNotBelowMissing(t *testing.T) { + // A signed vector channel (cosine/dot) can retrieve a document with negative + // similarity. Under the union's missing-uid=0 convention, that document must not + // fuse BELOW a document the channel never retrieved (which contributes 0). uid3 is + // retrieved by vec with a negative cosine; uid2 is absent from vec entirely. + text := fuseChannel{scores: map[uint64]float64{1: 10.0}, weight: 1.0} + vec := fuseChannel{scores: map[uint64]float64{1: 0.9, 3: -0.4}, weight: 1.0} + + res := fuseChannels([]fuseChannel{text, vec}, + fuseOpts{method: fusionLinear, normalize: normalizeMax}) + got := asMap(res) + + // uid3's negative similarity is clamped to 0, so it ties the missing baseline + // rather than going negative. + require.InDelta(t, 0.0, got[3], 1e-9) + // uid2 is absent from every channel, so it is not in the union at all. + require.NotContains(t, got, uint64(2)) + // A retrieved-but-dissimilar doc (uid3, score 0) is never ranked below a doc that + // would only appear via a negative contribution — no fused score is negative. + for uid, s := range got { + require.GreaterOrEqual(t, s, 0.0, "uid %d fused below the missing baseline", uid) + } +} + +func TestFuse_TopKTruncation(t *testing.T) { + a := ch(map[uint64]float64{1: 9, 2: 8, 3: 7, 4: 6, 5: 5}) + res := fuseChannels([]fuseChannel{a}, fuseOpts{method: fusionRRF, k: 60, topk: 3}) + require.Len(t, res, 3) + require.Equal(t, uint64(1), res[0].uid) + require.Equal(t, uint64(3), res[2].uid) +} + +func TestFuse_TopKHeapTieBreakEviction(t *testing.T) { + // Pins topKFused's tie-break eviction: with topk=2 and three uids tied at the + // same fused score, the kept pair must be the two LOWEST uids ({1, 2}), the + // same (score desc, uid asc) order the full-sort path produces. A score-only + // heap would keep whichever pair the map iteration happened to visit — this + // exercises the (cand.score == root.score && cand.uid < root.uid) replacement + // branch, which is NOT dead code. + a := ch(map[uint64]float64{1: 5.0, 2: 5.0, 3: 5.0}) + res := fuseChannels([]fuseChannel{a}, fuseOpts{method: fusionLinear, normalize: normalizeNone, topk: 2}) + require.Len(t, res, 2) + require.Equal(t, uint64(1), res[0].uid) + require.Equal(t, uint64(2), res[1].uid) + + // And the heap path must agree with the full-sort path on a mixed corpus. + b := ch(map[uint64]float64{7: 1.0, 8: 3.0, 9: 3.0, 10: 2.0, 11: 3.0}) + full := fuseChannels([]fuseChannel{b}, fuseOpts{method: fusionLinear, normalize: normalizeNone}) + topped := fuseChannels([]fuseChannel{b}, fuseOpts{method: fusionLinear, normalize: normalizeNone, topk: 3}) + require.Equal(t, full[:3], topped, "heap top-k must equal the full-sort prefix") +} + +func TestFuse_SingleChannelPassthroughOrder(t *testing.T) { + a := ch(map[uint64]float64{1: 1, 2: 9, 3: 5}) + res := fuseChannels([]fuseChannel{a}, fuseOpts{method: fusionRRF, k: 60}) + // Order should reflect channel ranking: 2 (rank1), 3 (rank2), 1 (rank3). + require.Equal(t, []uint64{2, 3, 1}, []uint64{res[0].uid, res[1].uid, res[2].uid}) +} + +func TestFuse_EmptyChannels(t *testing.T) { + res := fuseChannels([]fuseChannel{ch(nil), ch(map[uint64]float64{})}, + fuseOpts{method: fusionRRF, k: 60}) + require.Empty(t, res) +} + +func TestScoresFromVar_DropsNonFinite(t *testing.T) { + // Non-finite scores from a channel must be dropped so they can't break the sort + // comparator or poison linear sums. + m := types.NewShardedMap() + m.Set(1, types.Val{Tid: types.FloatID, Value: 0.5}) + m.Set(2, types.Val{Tid: types.FloatID, Value: math.NaN()}) + m.Set(3, types.Val{Tid: types.FloatID, Value: math.Inf(1)}) + m.Set(4, types.Val{Tid: types.FloatID, Value: math.Inf(-1)}) + m.Set(5, types.Val{Tid: types.FloatID, Value: 2.0}) + + scores, err := scoresFromVar(varValue{Vals: m}, "ch") + require.NoError(t, err) + require.Len(t, scores, 2, "only finite scores should survive") + require.Contains(t, scores, uint64(1)) + require.Contains(t, scores, uint64(5)) + require.NotContains(t, scores, uint64(2)) + require.NotContains(t, scores, uint64(3)) + require.NotContains(t, scores, uint64(4)) +} + +func TestFuseLinear_NonFiniteChannelDoesNotPoison(t *testing.T) { + // Even if a NaN slips into a channel passed directly to the core, max-normalize + // must not produce a NaN denominator that propagates. + bad := fuseChannel{scores: map[uint64]float64{1: math.NaN(), 2: math.NaN()}, weight: 1.0} + good := fuseChannel{scores: map[uint64]float64{1: 4.0}, weight: 1.0} + res := fuseChannels([]fuseChannel{bad, good}, fuseOpts{method: fusionLinear, normalize: normalizeMax}) + for _, r := range res { + require.False(t, math.IsNaN(r.score), "uid %d score must not be NaN", r.uid) + } +} + +func TestFuse_Determinism(t *testing.T) { + a := ch(map[uint64]float64{1: 5, 2: 5, 3: 5}) + b := ch(map[uint64]float64{3: 1, 2: 1, 1: 1}) + first := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionRRF, k: 60}) + for i := 0; i < 20; i++ { + again := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionRRF, k: 60}) + require.Equal(t, first, again) + } +} + +func TestFuseLinear_WindowDependentNormalization(t *testing.T) { + // SEMANTICS-OPEN (Q9): max-normalization divides by channelMaxAbs computed over + // whatever scores the channel var happened to retrieve (its window), not over any + // global corpus maximum — no global max exists anywhere in the pipeline. The SAME + // raw score for the SAME uid therefore normalizes differently depending on which + // other documents share its retrieval window (e.g. a bm25 channel with first:N vs + // unbounded). This test demonstrates and pins that window dependence; whether the + // denominator should instead be window-invariant is an open design decision. + const raw = 2.0 + + // Window 1: the channel also retrieved a doc scoring 10.0, so max is 10.0. + wide := fuseChannel{scores: map[uint64]float64{1: 10.0, 2: raw}, weight: 1.0} + res := fuseChannels([]fuseChannel{wide}, + fuseOpts{method: fusionLinear, normalize: normalizeMax}) + require.InDelta(t, 0.2, asMap(res)[2], 1e-9, "raw 2.0 under window max 10.0") + + // Window 2 (e.g. a deeper offset window): raw 2.0 is now the window max. + narrow := fuseChannel{scores: map[uint64]float64{2: raw, 3: 1.0}, weight: 1.0} + res = fuseChannels([]fuseChannel{narrow}, + fuseOpts{method: fusionLinear, normalize: normalizeMax}) + require.InDelta(t, 1.0, asMap(res)[2], 1e-9, "same raw 2.0 under window max 2.0") +} + +func TestFuseRRF_MassTiePlateau(t *testing.T) { + // SEMANTICS-OPEN (Q9): channelRanks assigns 1-based POSITIONAL ranks with a uid + // tie-break, so a plateau of identical scores receives ranks 1..n and RRF + // contributions ranging from 1/(k+1) down to 1/(k+n) — equal-relevance documents + // get very different fused scores purely by uid order. Dense ranking (all tied + // docs sharing one rank) is not implemented; this pins the positional choice + // until that decision is made. + const n = 1000 + scores := make(map[uint64]float64, n) + for uid := uint64(1); uid <= n; uid++ { + scores[uid] = 5.0 + } + res := fuseChannels([]fuseChannel{ch(scores)}, fuseOpts{method: fusionRRF, k: 60}) + require.Len(t, res, n) + got := asMap(res) + + const k = 60.0 + require.InDelta(t, 1/(k+1), got[1], 1e-12, "first tied uid gets rank 1") + require.InDelta(t, 1/(k+n), got[n], 1e-12, "last tied uid gets rank n") + // The spread across a single tie plateau: (k+n)/(k+1) = 1060/61 ≈ 17.4x. + require.Greater(t, got[1]/got[n], 10.0, + "positional ranking spreads identical-score docs by >10x under k=60") + // Output order across the plateau is uid ascending (rank order == uid order). + require.Equal(t, uint64(1), res[0].uid) + require.Equal(t, uint64(n), res[n-1].uid) +} + +func TestParseFuseOpts_NegativeWeights(t *testing.T) { + // SEMANTICS-OPEN (Q9): parseFuseOpts rejects only NaN/Inf weights — negative and + // zero weights are ACCEPTED today, even though fuseLinear's clamp-then-weight + // order means negative weights produce fused scores outside any documented + // range (see TestFuseLinear_NegativeWeightOutput). Whether to reject them is an + // open decision; this pins the current accept behavior. + args := func(spec string) []dql.Arg { + return []dql.Arg{{Value: "weights"}, {Value: spec}} + } + + _, weights, err := parseFuseOpts(args("-1,1"), 2) + require.NoError(t, err, "negative weight is currently accepted") + require.Equal(t, []float64{-1, 1}, weights) + + _, weights, err = parseFuseOpts(args("0,1"), 2) + require.NoError(t, err, "zero weight is currently accepted") + require.Equal(t, []float64{0, 1}, weights) + + // Non-finite weights are rejected (strconv.ParseFloat accepts these spellings, + // so the explicit IsNaN/IsInf guard is what fires). + for _, bad := range []string{"NaN,1", "Inf,1", "+Inf,1", "-Infinity,1"} { + _, _, err := parseFuseOpts(args(bad), 2) + require.Error(t, err, "weight spec %q must be rejected", bad) + require.Contains(t, err.Error(), "invalid weight") + } +} + +func TestFuseLinear_NegativeWeightOutput(t *testing.T) { + // SEMANTICS-OPEN (Q9): with an accepted negative weight, fuseLinear emits + // negative fused scores — outside the [0, sum(weights)] range that max-normalized + // non-negative weights would guarantee. Pin the exact arithmetic so any future + // rejection (or re-ranged) fix consciously changes this test. + a := fuseChannel{scores: map[uint64]float64{1: 4.0, 2: 2.0}, weight: -1.0} + b := fuseChannel{scores: map[uint64]float64{2: 3.0}, weight: 1.0} + res := fuseChannels([]fuseChannel{a, b}, + fuseOpts{method: fusionLinear, normalize: normalizeMax}) + got := asMap(res) + + // uid1: -1*(4/4) = -1.0 (a negative fused score escapes into the ranking). + require.InDelta(t, -1.0, got[1], 1e-9) + // uid2: -1*(2/4) + 1*(3/3) = 0.5 + require.InDelta(t, 0.5, got[2], 1e-9) + require.Equal(t, uint64(2), res[0].uid, "negatively-fused doc sinks below the positive one") +} + +func TestFuseLinear_ClampBeforeWeight(t *testing.T) { + // Pin (Q9): the negative-similarity clamp applies to the NORMALIZED score BEFORE + // the channel weight multiplies in. A negative raw score under a negative weight + // must therefore contribute exactly 0 (clamp(-0.5)=0, then -1*0=0) — never a + // positive contribution, which weighting-before-clamping would produce + // (-1 * -0.5 = +0.5). + c := fuseChannel{scores: map[uint64]float64{1: -0.5, 2: 1.0}, weight: -1.0} + res := fuseChannels([]fuseChannel{c}, + fuseOpts{method: fusionLinear, normalize: normalizeMax}) + got := asMap(res) + + require.InDelta(t, 0.0, got[1], 1e-9, "clamped-to-zero norm stays 0 under any weight") + require.InDelta(t, -1.0, got[2], 1e-9, "positive norm 1.0 weighted by -1") +} + +func TestFuseRRF_WeightSemanticsDeviateFromStandardRRF(t *testing.T) { + // INTENTIONAL DEVIATION (Q9): standard RRF (Cormack et al. 2009) is the + // unweighted sum over channels of 1/(k+rank). This implementation multiplies each + // channel's reciprocal-rank term by that channel's weight — weight * 1/(k+rank) — + // so that fuse(..., weights:) biases channels under method:rrf instead of + // silently ignoring the option. With the default weight 1.0 it reduces to + // standard RRF (see TestFuseRRF_DefaultWeightIsStandardRRF). Pin the exact + // weighted math on a case standard RRF would score as a dead tie. + a := fuseChannel{scores: map[uint64]float64{1: 9.0, 2: 5.0}, weight: 0.25} + b := fuseChannel{scores: map[uint64]float64{2: 9.0, 1: 5.0}, weight: 4.0} + res := fuseChannels([]fuseChannel{a, b}, fuseOpts{method: fusionRRF, k: 60}) + got := asMap(res) + + const k = 60.0 + // uid1: rank1 in a, rank2 in b. uid2: mirror image. + require.InDelta(t, 0.25/(k+1)+4.0/(k+2), got[1], 1e-12) + require.InDelta(t, 0.25/(k+2)+4.0/(k+1), got[2], 1e-12) + // Standard (unweighted) RRF would tie uid1 and uid2 exactly; the weights break + // the tie in favor of the heavier channel's top document. + require.Equal(t, uint64(2), res[0].uid) + require.Greater(t, got[2], got[1]) +} diff --git a/query/groupby.go b/query/groupby.go index 56bdbcd25df..ca42f4b6dfd 100644 --- a/query/groupby.go +++ b/query/groupby.go @@ -139,7 +139,10 @@ func aggregateGroup(grp *groupResult, child *SubGraph) (types.Val, error) { continue } - if len(child.valueMatrix[idx].Values) == 0 { + // valueMatrix can be shorter than SrcUIDs for coordinator-synthesized value + // variables (e.g. fused scores), which never populate a full matrix — guard + // the index or an aggregation over such a var panics. + if idx >= len(child.valueMatrix) || len(child.valueMatrix[idx].Values) == 0 { continue } v := child.valueMatrix[idx].Values[0] diff --git a/query/query.go b/query/query.go index 6926e2ac6ed..97f0f99d25e 100644 --- a/query/query.go +++ b/query/query.go @@ -7,6 +7,7 @@ package query import ( "context" + "encoding/binary" "fmt" "math" "sort" @@ -268,6 +269,15 @@ type SubGraph struct { // In graph terms, a list is a slice of outgoing edges from a node. uidMatrix []*pb.List + // rankerScores maps a matched document UID to its ranker score (BM25 relevance + // or vector similarity). It is snapshotted from the (uid-aligned) worker result + // the moment it arrives, before filters or pagination can shrink/reorder + // uidMatrix out of step with valueMatrix. populateUidValVar binds the score + // variable from this map keyed by UID, so the score stays correct even when the + // ranker block carries an @filter. nil unless the source function is a ranker + // (bm25 / similar_to). + rankerScores map[uint64]float64 + // facetsMatrix contains the facet values. There would a list corresponding to each uid in // uidMatrix. facetsMatrix []*pb.FacetsList @@ -1556,6 +1566,17 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su var ok bool switch { + case sg.SrcFunc != nil && sg.SrcFunc.Name == "fuse": + // Native hybrid search: fuse() combines several already-scored value + // variables (its NeedsVar channels) into one ranked value variable. Fusion is + // a coordinator-side operation over resolved variables — the channel blocks + // have already populated doneVars by the time this block is scheduled. We bind + // both the union uid set (uid(var)) and the uid->fused-score map (val(var)). + fv, err := computeFuse(sg.SrcFunc.Args, sg.Params.NeedsVar, doneVars, sgPath) + if err != nil { + return err + } + doneVars[sg.Params.Var] = fv case len(sg.counts) > 0: // 1. When count of a predicate is assigned a variable, we store the mapping of uid => // count(predicate). @@ -1591,6 +1612,27 @@ func (sg *SubGraph) populateUidValVar(doneVars map[string]varValue, sgPath []*Su Value: int64(len(sg.SrcUIDs.Uids)), } doneVars[sg.Params.Var].Vals.Set(math.MaxUint64, val) + case sg.SrcFunc != nil && (sg.SrcFunc.Name == "bm25" || sg.SrcFunc.Name == "similar_to") && + sg.rankerScores != nil: + // A query-side ranker (BM25 relevance or vector similarity) binds its + // per-document score as a value variable. We populate BOTH the matched uid set + // and the uid->score map so the variable works with uid(var), val(var) and + // orderdesc: val(var) — surfacing and ordering by score without a + // pseudo-predicate or a ParentVars channel. Scores are looked up from the + // uid-keyed snapshot taken at result time (sg.rankerScores), so they remain + // correct even after an @filter on the ranker block shrinks DestUIDs. For + // similar_to the score is a higher-is-better similarity; this also lets vector + // results feed fuse(). + if v, ok = doneVars[sg.Params.Var]; !ok { + v = varValue{Vals: types.NewShardedMap(), path: sgPath, strList: sg.valueMatrix} + } + v.Uids = sg.DestUIDs + for _, uid := range sg.DestUIDs.GetUids() { + if score, has := sg.rankerScores[uid]; has { + v.Vals.Set(uid, types.Val{Tid: types.FloatID, Value: score}) + } + } + doneVars[sg.Params.Var] = v case len(sg.DestUIDs.Uids) != 0 || (sg.Attr == "uid" && sg.SrcUIDs != nil): // 3. A uid variable. The variable could be defined in one of two places. // a) Either on the actual predicate. @@ -2173,6 +2215,7 @@ func ProcessGraph(ctx context.Context, sg, parent *SubGraph, rch chan error) { rch <- nil return } + var err error switch { case parent == nil && sg.SrcFunc != nil && sg.SrcFunc.Name == "uid": @@ -2275,6 +2318,36 @@ func ProcessGraph(ctx context.Context, sg, parent *SubGraph, rch chan error) { sg.List = result.List sg.vectorMetrics = result.VectorMetrics + // bm25 and similar_to return their per-document scores in valueMatrix + // positionally aligned with uidMatrix[0]. Snapshot them into a uid-keyed + // map now, while the two are still aligned — later filters/pagination + // shrink uidMatrix without touching valueMatrix, which would otherwise + // misbind scores to UIDs (and feed wrong scores into fuse() channels). + if sg.SrcFunc != nil && (sg.SrcFunc.Name == "bm25" || sg.SrcFunc.Name == "similar_to") && + len(result.UidMatrix) > 0 { + uids := result.UidMatrix[0].GetUids() + sg.rankerScores = make(map[uint64]float64, len(uids)) + for idx, uid := range uids { + if idx >= len(result.ValueMatrix) || len(result.ValueMatrix[idx].Values) == 0 { + continue + } + tv := result.ValueMatrix[idx].Values[0] + if len(tv.Val) != 8 { + continue + } + score := math.Float64frombits(binary.LittleEndian.Uint64(tv.Val)) + // Drop non-finite scores at the source: a NaN here would reach the + // value-var sort comparator (undefined order) and val() output. + // fuse() already drops them per-channel; direct consumers must be + // protected the same way. The uid stays matched — it just carries + // no score (same as fuse's convention). + if math.IsNaN(score) || math.IsInf(score, 0) { + continue + } + sg.rankerScores[uid] = score + } + } + if sg.Params.DoCount { if len(sg.Filters) == 0 { // If there is a filter, we need to do more work to get the actual count. @@ -2373,9 +2446,12 @@ func ProcessGraph(ctx context.Context, sg, parent *SubGraph, rch chan error) { } if len(sg.Params.Order) == 0 && len(sg.Params.FacetsOrder) == 0 { - // for `has` function when there is no filtering and ordering, we fetch - // correct paginated results so no need to apply pagination here. - if !(len(sg.Filters) == 0 && sg.SrcFunc != nil && sg.SrcFunc.Name == "has") { + // For `has` and `bm25`, the worker already returns correctly paginated + // results (bm25 paginates over score order, which the uid-sorted query-layer + // pagination cannot reproduce), so applying pagination again here would + // double-apply first/offset. Skip it when there is no filtering/ordering. + if !(len(sg.Filters) == 0 && sg.SrcFunc != nil && + (sg.SrcFunc.Name == "has" || sg.SrcFunc.Name == "bm25")) { // There is no ordering. Just apply pagination and return. if err = sg.applyPagination(ctx); err != nil { rch <- err @@ -2452,6 +2528,7 @@ func ProcessGraph(ctx context.Context, sg, parent *SubGraph, rch chan error) { } child.SrcUIDs = sg.DestUIDs // Make the connection. + if child.IsInternal() { // We dont have to execute these nodes. continue @@ -2496,16 +2573,56 @@ func (sg *SubGraph) applyPagination(ctx context.Context) error { } sg.updateUidMatrix() - for i := range sg.uidMatrix { - // Apply the offsets. - start, end := x.PageRange(sg.Params.Count, sg.Params.Offset, len(sg.uidMatrix[i].Uids)) - sg.uidMatrix[i].Uids = sg.uidMatrix[i].Uids[start:end] + if sg.rankerScores != nil && sg.Params.Count >= 0 && sg.SrcFunc != nil && + (sg.SrcFunc.Name == "bm25" || sg.SrcFunc.Name == "similar_to") { + // A ranker root's page is defined by score order. The worker paginates by + // score itself when there is no @filter (First is only pushed down then); + // with a filter, pagination happens here after filtering — and slicing the + // uid-ascending matrix would silently return the lowest-uid page instead of + // the top-scored one. Select the [offset, offset+count) window by + // (score desc, uid asc) from the snapshot, then restore ascending uid order + // for the pipeline. Presentation order still comes from orderdesc: val(var). + for i := range sg.uidMatrix { + sg.uidMatrix[i].Uids = pageByScore(sg.uidMatrix[i].Uids, sg.rankerScores, + sg.Params.Count, sg.Params.Offset) + } + } else { + for i := range sg.uidMatrix { + // Apply the offsets. + start, end := x.PageRange(sg.Params.Count, sg.Params.Offset, len(sg.uidMatrix[i].Uids)) + sg.uidMatrix[i].Uids = sg.uidMatrix[i].Uids[start:end] + } } // Re-merge the UID matrix. sg.DestUIDs = algo.MergeSorted(sg.uidMatrix) return nil } +// pageByScore returns the [offset, offset+count) window of uids ranked by +// (score desc, uid asc), re-sorted to ascending uid order. count <= 0 means no +// limit. Uids missing from scores (cannot happen for a ranker root, where every +// matched uid is snapshotted) rank last at score 0. +func pageByScore(uids []uint64, scores map[uint64]float64, count, offset int) []uint64 { + ranked := make([]uint64, len(uids)) + copy(ranked, uids) + sort.Slice(ranked, func(i, j int) bool { + si, sj := scores[ranked[i]], scores[ranked[j]] + if si != sj { + return si > sj + } + return ranked[i] < ranked[j] + }) + if offset > len(ranked) { + offset = len(ranked) + } + ranked = ranked[offset:] + if count > 0 && count < len(ranked) { + ranked = ranked[:count] + } + sort.Slice(ranked, func(i, j int) bool { return ranked[i] < ranked[j] }) + return ranked +} + // applyOrderAndPagination orders each posting list by a given attribute // before applying pagination. func (sg *SubGraph) applyOrderAndPagination(ctx context.Context) error { @@ -2751,7 +2868,8 @@ func isValidArg(a string) bool { func isValidFuncName(f string) bool { switch f { case "anyofterms", "allofterms", "val", "regexp", "anyoftext", "alloftext", "ngram", - "has", "uid", "uid_in", "anyof", "allof", "type", "match", "similar_to": + "has", "uid", "uid_in", "anyof", "allof", "type", "match", "similar_to", "bm25", + "fuse": return true } return isInequalityFn(f) || types.IsGeoFunc(f) @@ -2959,6 +3077,15 @@ func (req *Request) ProcessQuery(ctx context.Context) (err error) { continue } + // fuse() is a coordinator-side fusion over already-resolved value + // variables; it is never dispatched to a worker. The fused variable is + // produced from doneVars in populateVarMap (the fuse case in + // populateUidValVar) once its channel variables are populated. + if sg.SrcFunc != nil && sg.SrcFunc.Name == "fuse" { + errChan <- nil + continue + } + switch { case sg.Params.Alias == "shortest": // We allow only one shortest path block per query. diff --git a/query/query_bm25_test.go b/query/query_bm25_test.go new file mode 100644 index 00000000000..a70fa1e1976 --- /dev/null +++ b/query/query_bm25_test.go @@ -0,0 +1,1534 @@ +//go:build integration || cloud + +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +//nolint:lll +package query + +import ( + "context" + "encoding/json" + "fmt" + "math" + "strconv" + "strings" + "sync" + "testing" + + "github.com/dgraph-io/dgo/v250/protos/api" + "github.com/stretchr/testify/require" +) + +// uidHex queries Dgraph for the hex UID string of a given decimal UID. +// This avoids hardcoding hex values that depend on UID assignment order. +func uidHex(t *testing.T, decimalUID int) string { + t.Helper() + js := processQueryNoErr(t, fmt.Sprintf(`{ me(func: uid(%d)) { uid } }`, decimalUID)) + var resp struct { + Data struct { + Me []struct { + UID string `json:"uid"` + } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + require.NotEmpty(t, resp.Data.Me, "UID %d should exist", decimalUID) + return resp.Data.Me[0].UID +} + +func TestBM25Basic(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "quick brown fox")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + // Should return documents containing "quick", "brown", or "fox" + require.Contains(t, js, "quick brown fox jumps") + require.Contains(t, js, "quick brown fox leaps") +} + +func TestBM25Ordering(t *testing.T) { + // BM25 returns all matching documents. Use first:1 to verify the highest-scored + // document is "fox fox fox" (tf=3, short doc). + query := ` + { + me(func: bm25(description_bm25, "fox")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + // Should contain all fox-mentioning documents. + require.Contains(t, js, "fox fox fox") + require.Contains(t, js, "quick brown fox jumps") + + // first:1 should return the top-ranked document. + topQuery := ` + { + me(func: bm25(description_bm25, "fox"), first: 1) { + uid + description_bm25 + } + } + ` + topJs := processQueryNoErr(t, topQuery) + require.Contains(t, topJs, "fox fox fox", + "top-1 BM25 result for 'fox' should be 'fox fox fox' (highest tf, shortest doc)") +} + +func TestBM25WithParams(t *testing.T) { + // Custom k and b parameters + query := ` + { + me(func: bm25(description_bm25, "fox", "1.5", "0.5")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + require.Contains(t, js, "fox") +} + +func TestBM25InvalidParams(t *testing.T) { + // Negative k should be rejected. + query := ` + { + me(func: bm25(description_bm25, "fox", "-1.0", "0.75")) { + uid + } + } + ` + _, err := processQuery(context.Background(), t, query) + require.Error(t, err) + require.Contains(t, err.Error(), "bm25: k must be a positive finite number") + + // b > 1 should be rejected. + query2 := ` + { + me(func: bm25(description_bm25, "fox", "1.2", "1.5")) { + uid + } + } + ` + _, err = processQuery(context.Background(), t, query2) + require.Error(t, err) + require.Contains(t, err.Error(), "bm25: b must be between 0 and 1") + + // b < 0 should be rejected. + query3 := ` + { + me(func: bm25(description_bm25, "fox", "1.2", "-0.5")) { + uid + } + } + ` + _, err = processQuery(context.Background(), t, query3) + require.Error(t, err) + require.Contains(t, err.Error(), "bm25: b must be between 0 and 1") +} + +func TestBM25AsFilter(t *testing.T) { + query := ` + { + me(func: has(description_bm25)) @filter(bm25(description_bm25, "fox")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + require.Contains(t, js, "fox") + // Should not contain documents without "fox" + require.NotContains(t, js, "Dogs are loyal") +} + +func TestBM25NoResults(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "xyznonexistent")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + require.JSONEq(t, `{"data": {"me":[]}}`, js) +} + +func TestBM25SingleTerm(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "dog")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + require.Contains(t, js, "dog") +} + +func TestBM25MultiTerm(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "quick lazy")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + // Should find docs with "quick" or "lazy" (scores accumulate). + // Doc 501 has both "quick" and "lazy", so it should rank high. + require.Contains(t, js, "quick brown fox jumps over the lazy dog") +} + +func TestBM25AllStopwords(t *testing.T) { + // A query consisting entirely of stopwords should return no results. + query := ` + { + me(func: bm25(description_bm25, "the a an")) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + require.JSONEq(t, `{"data": {"me":[]}}`, js) +} + +func TestBM25EmptyPredicate(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "")) { + uid + } + } + ` + js := processQueryNoErr(t, query) + require.JSONEq(t, `{"data": {"me":[]}}`, js) +} + +func TestBM25WithCount(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "fox")) { + count(uid) + } + } + ` + js := processQueryNoErr(t, query) + // Should have at least 2 results (docs with "fox") + require.Contains(t, js, "count") +} + +func TestBM25Pagination(t *testing.T) { + query := ` + { + me(func: bm25(description_bm25, "fox"), first: 1) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + // With first:1, should return exactly one result (the highest-scoring). + // Doc 503 "fox fox fox" should be the top result. + require.Contains(t, js, "fox fox fox") +} + +func TestBM25ScoreOrdering(t *testing.T) { + // Bind the bm25 score to a value variable and order results by it via val(). + query := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score), first: 1) { + uid + description_bm25 + val(score) + } + } + ` + js := processQueryNoErr(t, query) + // "fox fox fox" (doc 503) has the highest BM25 score (tf=3, shortest doc). + require.Contains(t, js, "fox fox fox") +} + +func TestBM25ScoreOrderingMultiTerm(t *testing.T) { + // Multi-term query with score ordering: "quick lazy" should rank doc 501 highest + // since it contains both terms. + query := ` + { + score as var(func: bm25(description_bm25, "quick lazy")) + me(func: uid(score), orderdesc: val(score), first: 1) { + uid + description_bm25 + val(score) + } + } + ` + js := processQueryNoErr(t, query) + require.Contains(t, js, "quick brown fox jumps over the lazy dog") +} + +func TestBM25ScoreOrderingAllResults(t *testing.T) { + // Verify all results are returned in score-descending order via val(score). + query := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score)) { + uid + description_bm25 + val(score) + } + } + ` + js := processQueryNoErr(t, query) + // All fox-containing docs should appear. + require.Contains(t, js, "fox fox fox") + require.Contains(t, js, "quick brown fox jumps") + // Score values should be present. + require.Contains(t, js, "val(score)") +} + +func TestBM25ScoreWithPagination(t *testing.T) { + // Use offset with score ordering. + query := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score), first: 1, offset: 1) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + // Should return the second-highest scored document (not "fox fox fox"). + require.NotContains(t, js, "fox fox fox") + require.Contains(t, js, "fox") +} + +// parseScoresFromJSON extracts uid → score from JSON responses containing val(score). +func parseScoresFromJSON(t *testing.T, js string) map[string]float64 { + t.Helper() + var resp struct { + Data struct { + Me []struct { + UID string `json:"uid"` + Score float64 `json:"val(score)"` + } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + scores := make(map[string]float64) + for _, item := range resp.Data.Me { + scores[item.UID] = item.Score + } + return scores +} + +func TestBM25IncrementalAddBatch(t *testing.T) { + batch1 := ` + <600> "alpha bravo charlie" . + <601> "delta echo foxtrot" . + ` + batch2 := ` + <602> "golf hotel india" . + <603> "juliet kilo lima" . + <604> "mike november oscar" . + ` + batch3 := ` + <605> "papa quebec romeo" . + <606> "sierra tango uniform" . + <607> "victor whiskey xray" . + ` + cleanup := func() { + deleteTriplesInCluster(` + <600> * . + <601> * . + <602> * . + <603> * . + <604> * . + <605> * . + <606> * . + <607> * . + `) + } + t.Cleanup(cleanup) + + countQuery := ` + { + me(func: bm25(description_bm25, "alpha bravo delta echo golf juliet mike papa sierra victor")) { + count(uid) + } + } + ` + + // Batch 1: add 2 docs. + require.NoError(t, addTriplesToCluster(batch1)) + js := processQueryNoErr(t, countQuery) + require.Contains(t, js, `"count":2`) + + // Batch 2: add 3 more docs → total 5. + require.NoError(t, addTriplesToCluster(batch2)) + js = processQueryNoErr(t, countQuery) + require.Contains(t, js, `"count":5`) + + // Batch 3: add 3 more docs → total 8. + require.NoError(t, addTriplesToCluster(batch3)) + js = processQueryNoErr(t, countQuery) + require.Contains(t, js, `"count":8`) + + // Verify specific new terms are searchable. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "whiskey")) { uid description_bm25 } }`) + require.Contains(t, js, "whiskey") +} + +func TestBM25CorpusStatsAffectIDF(t *testing.T) { + // Capture baseline score for "fox" query. + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + } + ` + jsBefore := processQueryNoErr(t, scoreQuery) + scoresBefore := parseScoresFromJSON(t, jsBefore) + require.NotEmpty(t, scoresBefore, "baseline should have fox results") + + // Add 10 non-fox docs → N grows, df("fox") stays same → IDF should increase. + var triples string + for i := 610; i < 620; i++ { + triples += fmt.Sprintf(`<%d> "completely unrelated document about cats and dogs number %d" . +`, i, i) + } + require.NoError(t, addTriplesToCluster(triples)) + t.Cleanup(func() { + var del string + for i := 610; i < 620; i++ { + del += fmt.Sprintf("<%d> * .\n", i) + } + deleteTriplesInCluster(del) + }) + + jsAfter := processQueryNoErr(t, scoreQuery) + scoresAfter := parseScoresFromJSON(t, jsAfter) + + // Compare score for UID 503 ("fox fox fox") — should increase. + uid503 := uidHex(t, 503) + before, ok1 := scoresBefore[uid503] + after, ok2 := scoresAfter[uid503] + require.True(t, ok1 && ok2, "UID 503 should appear in both before and after results") + require.Greater(t, after, before, + "IDF should increase when corpus grows with non-matching docs (before=%f, after=%f)", before, after) +} + +func TestBM25DocumentUpdate(t *testing.T) { + // Add a doc with lots of "fox". + require.NoError(t, addTriplesToCluster(`<620> "fox fox fox fox" .`)) + t.Cleanup(func() { + deleteTriplesInCluster(`<620> * .`) + }) + + uid620 := uidHex(t, 620) + + // Should rank top for "fox". + js := processQueryNoErr(t, ` + { + me(func: bm25(description_bm25, "fox"), first: 1) { + uid + } + }`) + require.Contains(t, js, `"`+uid620+`"`) + + // Update to remove "fox", add "cat". + deleteTriplesInCluster(`<620> "fox fox fox fox" .`) + require.NoError(t, addTriplesToCluster(`<620> "the cat sat on the mat" .`)) + + // Should no longer appear in "fox" results. + js = processQueryNoErr(t, ` + { + me(func: bm25(description_bm25, "fox")) { + uid + } + }`) + require.NotContains(t, js, `"`+uid620+`"`) + + // Should appear in "cat" results. + js = processQueryNoErr(t, ` + { + me(func: bm25(description_bm25, "cat")) { + uid + } + }`) + require.Contains(t, js, `"`+uid620+`"`) +} + +func TestBM25DocumentDeletion(t *testing.T) { + require.NoError(t, addTriplesToCluster(`<625> "unique elephant term" .`)) + t.Cleanup(func() { + // Cleanup in case test fails before explicit delete. + deleteTriplesInCluster(`<625> * .`) + }) + + uid625 := uidHex(t, 625) + + // Should find the elephant doc. + js := processQueryNoErr(t, `{ me(func: bm25(description_bm25, "elephant")) { uid } }`) + require.Contains(t, js, `"`+uid625+`"`) + + // Delete it. + deleteTriplesInCluster(`<625> "unique elephant term" .`) + + // Should return empty for "elephant". + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "elephant")) { uid } }`) + require.JSONEq(t, `{"data": {"me":[]}}`, js) + + // Baseline "fox" results should be unaffected. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "fox")) { uid description_bm25 } }`) + require.Contains(t, js, "fox") +} + +func TestBM25ScoreStabilityAsCorpusGrows(t *testing.T) { + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + } + ` + uid503 := uidHex(t, 503) + + // Phase 1: baseline score. + js1 := processQueryNoErr(t, scoreQuery) + scores1 := parseScoresFromJSON(t, js1) + score1, ok := scores1[uid503] + require.True(t, ok, "UID 503 must appear in baseline") + + // Phase 2: add 5 fox docs → IDF decreases. + var foxTriples string + for i := 630; i < 635; i++ { + foxTriples += fmt.Sprintf(`<%d> "the fox runs quickly across the field number %d" . +`, i, i) + } + require.NoError(t, addTriplesToCluster(foxTriples)) + t.Cleanup(func() { + var del string + for i := 630; i < 640; i++ { + del += fmt.Sprintf("<%d> * .\n", i) + } + deleteTriplesInCluster(del) + }) + + js2 := processQueryNoErr(t, scoreQuery) + scores2 := parseScoresFromJSON(t, js2) + score2, ok := scores2[uid503] + require.True(t, ok, "UID 503 must appear after adding fox docs") + require.Greater(t, score1, score2, + "Adding fox docs should decrease IDF and thus score (phase1=%f, phase2=%f)", score1, score2) + + // Phase 3: add 5 non-fox docs → IDF increases relative to phase 2. + var nonFoxTriples string + for i := 635; i < 640; i++ { + nonFoxTriples += fmt.Sprintf(`<%d> "unrelated content about birds and fish number %d" . +`, i, i) + } + require.NoError(t, addTriplesToCluster(nonFoxTriples)) + + js3 := processQueryNoErr(t, scoreQuery) + scores3 := parseScoresFromJSON(t, js3) + score3, ok := scores3[uid503] + require.True(t, ok, "UID 503 must appear after adding non-fox docs") + require.Greater(t, score3, score2, + "Adding non-fox docs should increase IDF relative to phase2 (phase2=%f, phase3=%f)", score2, score3) +} + +func TestBM25LargeCorpus(t *testing.T) { + // Add 100 docs: 50 with "alpha", 50 with "beta". + var triples string + for i := 700; i < 750; i++ { + triples += fmt.Sprintf(`<%d> "alpha document content number %d with some padding words" . +`, i, i) + } + for i := 750; i < 800; i++ { + triples += fmt.Sprintf(`<%d> "beta document content number %d with some padding words" . +`, i, i) + } + require.NoError(t, addTriplesToCluster(triples)) + t.Cleanup(func() { + var del string + for i := 700; i < 800; i++ { + del += fmt.Sprintf("<%d> * .\n", i) + } + deleteTriplesInCluster(del) + }) + + // Count alpha docs. + js := processQueryNoErr(t, `{ me(func: bm25(description_bm25, "alpha")) { count(uid) } }`) + require.Contains(t, js, `"count":50`) + + // Count beta docs. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "beta")) { count(uid) } }`) + require.Contains(t, js, `"count":50`) + + // Union count: "alpha beta" should match all 100. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "alpha beta")) { count(uid) } }`) + require.Contains(t, js, `"count":100`) + + // Pagination: first:10, offset:40 for alpha should return 10 results. + js = processQueryNoErr(t, ` + { + score as var(func: bm25(description_bm25, "alpha")) + me(func: uid(score), orderdesc: val(score), first: 10, offset: 40) { + uid + } + }`) + var resp struct { + Data struct { + Me []struct{ UID string } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + require.Len(t, resp.Data.Me, 10, "pagination first:10 offset:40 should return exactly 10 results") +} + +func TestBM25EdgeCaseSingleCharTerm(t *testing.T) { + require.NoError(t, addTriplesToCluster(`<640> "x y z" .`)) + t.Cleanup(func() { + deleteTriplesInCluster(`<640> * .`) + }) + + // Single-char terms may or may not be indexed depending on tokenizer. + // Just verify no panic/error. + _, err := processQuery(context.Background(), t, ` + { + me(func: bm25(description_bm25, "x")) { + uid + } + }`) + require.NoError(t, err) +} + +func TestBM25EdgeCaseLongDocument(t *testing.T) { + // Build a ~500-word document with "fox" appearing once. + words := make([]string, 500) + for i := range words { + words[i] = "padding" + } + words[250] = "fox" + longDoc := strings.Join(words, " ") + + require.NoError(t, addTriplesToCluster(fmt.Sprintf(`<645> %q .`, longDoc))) + t.Cleanup(func() { + deleteTriplesInCluster(`<645> * .`) + }) + + // Get scores for "fox" query. + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + } + ` + js := processQueryNoErr(t, scoreQuery) + scores := parseScoresFromJSON(t, js) + + uid503 := uidHex(t, 503) // "fox fox fox" (doclen=3) + uid645 := uidHex(t, 645) // long doc (doclen~500) + s503, ok1 := scores[uid503] + s645, ok2 := scores[uid645] + require.True(t, ok1, "UID 503 must appear in fox results") + require.True(t, ok2, "UID 645 must appear in fox results") + require.Greater(t, s503, s645, + "Short doc with high tf should score higher than long doc with low tf (503=%f, 645=%f)", s503, s645) +} + +func TestBM25EdgeCaseUnicode(t *testing.T) { + triples := ` + <650> "der schnelle braune Fuchs springt" . + <651> "le renard brun rapide saute" . + <652> "el zorro marrón rápido salta" . + ` + require.NoError(t, addTriplesToCluster(triples)) + t.Cleanup(func() { + deleteTriplesInCluster(` + <650> * . + <651> * . + <652> * . + `) + }) + + uid650 := uidHex(t, 650) + uid651 := uidHex(t, 651) + uid652 := uidHex(t, 652) + + // Query German term. + js := processQueryNoErr(t, `{ me(func: bm25(description_bm25, "Fuchs")) { uid } }`) + require.Contains(t, js, `"`+uid650+`"`) + + // Query French term. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "renard")) { uid } }`) + require.Contains(t, js, `"`+uid651+`"`) + + // Query Spanish term. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "zorro")) { uid } }`) + require.Contains(t, js, `"`+uid652+`"`) +} + +func TestBM25EdgeCaseAllStopwordsDoc(t *testing.T) { + require.NoError(t, addTriplesToCluster(`<655> "the a an is are was were" .`)) + t.Cleanup(func() { + deleteTriplesInCluster(`<655> * .`) + }) + + uid655 := uidHex(t, 655) + + // Query "the" — should return empty since "the" is a stopword. + js := processQueryNoErr(t, `{ me(func: bm25(description_bm25, "the")) { uid } }`) + require.NotContains(t, js, `"`+uid655+`"`) // 655 should not appear + + // But the doc should exist via has(). + js = processQueryNoErr(t, ` + { + me(func: has(description_bm25)) @filter(uid(655)) { + uid + } + }`) + require.Contains(t, js, `"`+uid655+`"`) +} + +func TestBM25WithUidFilter(t *testing.T) { + // BM25 root with uid filter to restrict results. + query := ` + { + me(func: bm25(description_bm25, "fox")) @filter(uid(501, 503)) { + uid + description_bm25 + } + } + ` + js := processQueryNoErr(t, query) + uid501 := uidHex(t, 501) + uid502 := uidHex(t, 502) + uid503 := uidHex(t, 503) + uid506 := uidHex(t, 506) + // Should contain only UIDs 501 and 503. + require.Contains(t, js, `"`+uid501+`"`) + require.Contains(t, js, `"`+uid503+`"`) + // Should NOT contain other fox docs like 502, 506. + require.NotContains(t, js, `"`+uid502+`"`) + require.NotContains(t, js, `"`+uid506+`"`) +} + +func TestBM25ScoreValuesAreValidFloats(t *testing.T) { + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "fox")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + } + ` + js := processQueryNoErr(t, scoreQuery) + scores := parseScoresFromJSON(t, js) + require.NotEmpty(t, scores, "should have at least one result") + + var prevScore float64 + first := true + // Iterate over results in order (they're orderdesc by score). + var resp struct { + Data struct { + Me []struct { + UID string `json:"uid"` + Score float64 `json:"val(score)"` + } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + + for _, item := range resp.Data.Me { + score := item.Score + require.False(t, math.IsNaN(score), "score should not be NaN for uid %s", item.UID) + require.False(t, math.IsInf(score, 0), "score should not be Inf for uid %s", item.UID) + require.Greater(t, score, 0.0, "score should be positive for uid %s", item.UID) + + if !first { + require.GreaterOrEqual(t, prevScore, score, + "scores should be in descending order: %f >= %f", prevScore, score) + } + prevScore = score + first = false + } +} + +func TestBM25IncrementalAddThenDeleteThenReadd(t *testing.T) { + t.Cleanup(func() { + deleteTriplesInCluster(`<670> * .`) + }) + + // Phase 1: add with "elephant". + require.NoError(t, addTriplesToCluster(`<670> "elephant roams the savanna" .`)) + uid670 := uidHex(t, 670) + js := processQueryNoErr(t, `{ me(func: bm25(description_bm25, "elephant")) { uid } }`) + require.Contains(t, js, `"`+uid670+`"`) + + // Phase 2: delete. + deleteTriplesInCluster(`<670> "elephant roams the savanna" .`) + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "elephant")) { uid } }`) + require.NotContains(t, js, `"`+uid670+`"`) + + // Phase 3: re-add with different content. + require.NoError(t, addTriplesToCluster(`<670> "penguin waddles on the ice" .`)) + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "penguin")) { uid } }`) + require.Contains(t, js, `"`+uid670+`"`) + + // "elephant" should still not match 670. + js = processQueryNoErr(t, `{ me(func: bm25(description_bm25, "elephant")) { uid } }`) + require.NotContains(t, js, `"`+uid670+`"`) +} + +func TestBM25NonIndexedPredicateError(t *testing.T) { + // "name" predicate does not have @index(bm25). + query := ` + { + me(func: bm25(name, "alice")) { + uid + } + } + ` + _, err := processQuery(context.Background(), t, query) + require.Error(t, err) + require.Contains(t, err.Error(), "bm25") +} + +func TestBM25ConcurrentBatchAdd(t *testing.T) { + // Add 5 batches of 4 docs each (UIDs 680-699) back-to-back. + t.Cleanup(func() { + var del string + for i := 680; i < 700; i++ { + del += fmt.Sprintf("<%d> * .\n", i) + } + deleteTriplesInCluster(del) + }) + + for batch := 0; batch < 5; batch++ { + var triples string + for j := 0; j < 4; j++ { + uid := 680 + batch*4 + j + triples += fmt.Sprintf(`<%d> "searchterm batch%d doc%d content here" . +`, uid, batch, j) + } + require.NoError(t, addTriplesToCluster(triples)) + } + + // All 20 docs should be findable. + js := processQueryNoErr(t, `{ me(func: bm25(description_bm25, "searchterm")) { count(uid) } }`) + require.Contains(t, js, `"count":20`) + + // Spot-check a doc from each batch. + for batch := 0; batch < 5; batch++ { + decUID := 680 + batch*4 + hexUID := uidHex(t, decUID) + term := fmt.Sprintf("batch%d", batch) + js = processQueryNoErr(t, fmt.Sprintf(`{ me(func: bm25(description_bm25, "%s")) { uid } }`, term)) + require.Contains(t, js, `"`+hexUID+`"`, "doc %d from batch %d should be searchable", decUID, batch) + } +} + +// parseCorpusCount returns the total number of documents with the description_bm25 predicate. +func parseCorpusCount(t *testing.T) float64 { + t.Helper() + js := processQueryNoErr(t, `{ me(func: has(description_bm25)) { count(uid) } }`) + var resp struct { + Data struct { + Me []struct { + Count int `json:"count"` + } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + require.NotEmpty(t, resp.Data.Me) + n := float64(resp.Data.Me[0].Count) + require.Greater(t, n, 0.0, "corpus must have documents") + return n +} + +func TestBM25ExactScoreValues(t *testing.T) { + // Exact score verification using b=0 (BM15 variant) to eliminate avgDL dependency. + // With b=0: score = idf * (k+1) * tf / (k + tf) + // This validates the core BM25 formula computes correct numerical values. + triples := ` + <850> "quasar quasar quasar" . + <851> "quasar nebula pulsar" . + ` + require.NoError(t, addTriplesToCluster(triples)) + t.Cleanup(func() { + deleteTriplesInCluster(` + <850> * . + <851> * . + `) + }) + + N := parseCorpusCount(t) + + // Query "quasar" with b=0 so score depends only on tf, k, and IDF (not avgDL). + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "quasar", "1.2", "0")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }` + js := processQueryNoErr(t, scoreQuery) + scores := parseScoresFromJSON(t, js) + + k := 1.2 + df := 2.0 // both 850 and 851 contain "quasar" + idf := math.Log1p((N - df + 0.5) / (df + 0.5)) + + // Doc 850 "quasar quasar quasar": tf=3, b=0 → score = idf * 2.2 * 3 / 4.2 + expected850 := idf * (k + 1) * 3.0 / (k + 3.0) + // Doc 851 "quasar nebula pulsar": tf=1, b=0 → score = idf * 2.2 * 1 / 2.2 = idf + expected851 := idf * (k + 1) * 1.0 / (k + 1.0) + + uid850 := uidHex(t, 850) + uid851 := uidHex(t, 851) + actual850, ok := scores[uid850] + require.True(t, ok, "UID 850 (%s) must be in results", uid850) + actual851, ok := scores[uid851] + require.True(t, ok, "UID 851 (%s) must be in results", uid851) + + require.InEpsilon(t, expected850, actual850, 1e-6, + "Doc 850 score mismatch: expected %f, got %f (N=%f, df=%f, idf=%f)", + expected850, actual850, N, df, idf) + require.InEpsilon(t, expected851, actual851, 1e-6, + "Doc 851 score mismatch: expected %f, got %f (N=%f, df=%f, idf=%f)", + expected851, actual851, N, df, idf) + + // Verify ordering: higher tf should yield higher score. + require.Greater(t, actual850, actual851) +} + +func TestBM25BM15NoLengthNormalization(t *testing.T) { + // With b=0 (BM15 variant), document length should NOT affect the score. + // Two docs with the same term frequency but different lengths must score identically. + triples := ` + <860> "vortex" . + <861> "vortex alpha bravo charlie delta echo foxtrot golf hotel india" . + ` + require.NoError(t, addTriplesToCluster(triples)) + t.Cleanup(func() { + deleteTriplesInCluster(` + <860> * . + <861> * . + `) + }) + + // Query with b=0: length normalization disabled. + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "vortex", "1.2", "0")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }` + js := processQueryNoErr(t, scoreQuery) + scores := parseScoresFromJSON(t, js) + + uid860 := uidHex(t, 860) + uid861 := uidHex(t, 861) + score860, ok1 := scores[uid860] + score861, ok2 := scores[uid861] + require.True(t, ok1, "UID 860 must be in results") + require.True(t, ok2, "UID 861 must be in results") + + // With b=0 and same tf=1, scores must be equal regardless of document length. + require.InDelta(t, score860, score861, 1e-9, + "b=0 should disable length normalization: short doc score=%f, long doc score=%f", + score860, score861) + + // Now verify that with default b=0.75, the shorter doc scores higher. + scoreQueryDefault := ` + { + score as var(func: bm25(description_bm25, "vortex")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }` + js = processQueryNoErr(t, scoreQueryDefault) + scoresDefault := parseScoresFromJSON(t, js) + + defScore860, ok1 := scoresDefault[uid860] + defScore861, ok2 := scoresDefault[uid861] + require.True(t, ok1, "UID 860 must be in default results") + require.True(t, ok2, "UID 861 must be in default results") + require.Greater(t, defScore860, defScore861, + "With b=0.75, shorter doc (doclen=1) should score higher than longer doc (doclen=10)") +} + +func TestBM25SingleMatchingDocument(t *testing.T) { + // Edge case: a single document matching the query term (df=1). + // IDF should be high since the term is very rare. + triples := `<865> "aardvark" .` + require.NoError(t, addTriplesToCluster(triples)) + t.Cleanup(func() { + deleteTriplesInCluster(`<865> * .`) + }) + + N := parseCorpusCount(t) + + // Query with b=0 for exact verification. + scoreQuery := ` + { + score as var(func: bm25(description_bm25, "aardvark", "1.2", "0")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }` + js := processQueryNoErr(t, scoreQuery) + scores := parseScoresFromJSON(t, js) + + require.Len(t, scores, 1, "exactly one document should match 'aardvark'") + + uid865 := uidHex(t, 865) + actual, ok := scores[uid865] + require.True(t, ok, "UID 865 (%s) must be in results", uid865) + + // With df=1, tf=1, b=0, k=1.2: + // idf = log1p((N - 1 + 0.5) / (1 + 0.5)) = log1p((N - 0.5) / 1.5) + // score = idf * 2.2 * 1 / (1.2 + 1) = idf * 2.2 / 2.2 = idf + k := 1.2 + df := 1.0 + idf := math.Log1p((N - df + 0.5) / (df + 0.5)) + expected := idf * (k + 1) * 1.0 / (k + 1.0) // simplifies to idf + + require.InEpsilon(t, expected, actual, 1e-6, + "Single-doc score mismatch: expected %f, got %f (N=%f, idf=%f)", + expected, actual, N, idf) + require.Greater(t, actual, 0.0, "score must be positive") + require.False(t, math.IsInf(actual, 0), "score must be finite") +} + +// TestBM25RebuildOnExistingData loads documents BEFORE a BM25 index exists, then adds +// @index(bm25) to trigger an index rebuild over the existing data. The streaming +// rebuild aggregates corpus statistics across many threads; a naive per-thread +// read-modify-write counter would undercount N/avgDL. This verifies bm25 ranks +// exactly as if the documents had been indexed live. +func TestBM25RebuildOnExistingData(t *testing.T) { + pred := "bm25_rebuild" + t.Cleanup(func() { dropPredicate(pred) }) + + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <920> <%[1]s> "fox fox fox" . + <921> <%[1]s> "fox dog" . + <922> <%[1]s> "dog cat bird fish" . + `, pred))) + + // Add the index now: this rebuilds over the three existing documents. + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + query := fmt.Sprintf(` + { + score as var(func: bm25(%s, "fox")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }`, pred) + scores := parseScoresFromJSON(t, processQueryNoErr(t, query)) + + require.Len(t, scores, 2, "both documents containing 'fox' must match after rebuild") + require.Greater(t, scores[uidHex(t, 920)], scores[uidHex(t, 921)], + "the denser, shorter document must rank higher after rebuild") + require.Greater(t, scores[uidHex(t, 920)], 0.0, "rebuilt stats must yield a positive score") +} + +// TestBM25DropThenReaddNoDoubleCount verifies that dropping the BM25 index clears its +// corpus statistics so re-adding it rebuilds from scratch. If the stats survived the +// drop, the rebuild would accumulate on top of stale counters, inflating N/avgDL and +// shifting every score. +func TestBM25DropThenReaddNoDoubleCount(t *testing.T) { + pred := "bm25_dropreadd" + t.Cleanup(func() { dropPredicate(pred) }) + + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <930> <%[1]s> "alpha alpha beta" . + <931> <%[1]s> "alpha gamma" . + <932> <%[1]s> "beta gamma delta" . + `, pred))) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + query := fmt.Sprintf(` + { + score as var(func: bm25(%s, "alpha")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }`, pred) + before := parseScoresFromJSON(t, processQueryNoErr(t, query)) + require.NotEmpty(t, before) + + // Drop the index (keeping the data), then re-add it to force a rebuild. + setSchema(fmt.Sprintf("%s: string .", pred)) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + after := parseScoresFromJSON(t, processQueryNoErr(t, query)) + + require.Equal(t, len(before), len(after), "result set size must be stable across drop/re-add") + for uid, score := range before { + require.InEpsilon(t, score, after[uid], 1e-9, + "score for %s must be identical after drop/re-add (no stale-stats double counting)", uid) + } +} + +// TestBM25ConcurrentOverlappingTxns adds documents whose UIDs share BM25 stats buckets +// from many goroutines at once. Corpus stats are guarded by a value-independent +// conflict key, so overlapping transactions to the same bucket conflict and one +// retries rather than silently overwriting the other's contribution. Every document +// must end up indexed and searchable. +func TestBM25ConcurrentOverlappingTxns(t *testing.T) { + pred := "bm25_concurrent" + t.Cleanup(func() { dropPredicate(pred) }) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + ctx := context.Background() + // UIDs chosen so several collide on uid%32 (the stats bucket count). + uids := []int{1000, 1032, 1064, 1001, 1033, 1065, 1002, 1034} + + var wg sync.WaitGroup + for _, uid := range uids { + wg.Add(1) + go func(uid int) { + defer wg.Done() + // Retry on conflict, exactly as a client would. + for { + txn := client.NewTxn() + _, err := txn.Mutate(ctx, &api.Mutation{ + SetNquads: []byte(fmt.Sprintf( + `<%d> <%s> "concurrent indexing term doc%d" .`, uid, pred, uid)), + CommitNow: true, + }) + _ = txn.Discard(ctx) + if err == nil { + return + } + } + }(uid) + } + wg.Wait() + + js := processQueryNoErr(t, fmt.Sprintf( + `{ me(func: bm25(%s, "concurrent")) { count(uid) } }`, pred)) + require.Contains(t, js, fmt.Sprintf(`"count":%d`, len(uids)), + "all concurrently-indexed documents must be searchable (no lost stats updates)") +} + +// parseUIDScoreRows extracts (uid, val(score)) rows in response order, decoding the +// hex uid strings to uint64. The score map is keyed by uid; blocks without val(score) +// simply yield zero-valued scores that callers ignore. +func parseUIDScoreRows(t *testing.T, js string) ([]uint64, map[uint64]float64) { + t.Helper() + var resp struct { + Data struct { + Me []struct { + UID string `json:"uid"` + Score float64 `json:"val(score)"` + } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + uids := make([]uint64, 0, len(resp.Data.Me)) + scores := make(map[uint64]float64, len(resp.Data.Me)) + for _, item := range resp.Data.Me { + uid, err := strconv.ParseUint(strings.TrimPrefix(item.UID, "0x"), 16, 64) + require.NoError(t, err, "uid %q must be a hex uid", item.UID) + uids = append(uids, uid) + scores[uid] = item.Score + } + return uids, scores +} + +// setupBM25FilterCorpus indexes a small corpus where the top-scored document for the +// query term "fox" has the HIGHEST uid, so a uid-ordered truncation of the result set +// is distinguishable from a score-ordered one. Every doc also carries an int predicate +// for @filter(ge(...)): the two top-scored docs get negative values so ge(pop, 0) can +// drop the entire unfiltered top of the ranking. +// +// Score order for "fox" (avgDL = 12/5 = 2.4): 300006 (tf=3, dl=3) > 300004 (tf=2, +// dl=2) > 300001 (tf=1, dl=2) > 300002 (tf=1, dl=3). Uid order of the matches is the +// exact opposite of score order for the top two. +// ensureBM25UidLease extends zero's uid lease to cover the 300001-350013 range these +// tests hardcode (TestMain only leases 65536; AssignUids advances maxLeasedUID). +var bm25UidLeaseOnce sync.Once + +func ensureBM25UidLease(t *testing.T) { + t.Helper() + bm25UidLeaseOnce.Do(func() { + require.NoError(t, dc.AssignUids(client.Dgraph, 400000)) + }) +} + +func setupBM25FilterCorpus(t *testing.T, textPred, popPred string) { + t.Helper() + ensureBM25UidLease(t) + t.Cleanup(func() { + dropPredicate(textPred) + dropPredicate(popPred) + }) + setSchema(fmt.Sprintf("%s: string @index(bm25) .\n%s: int @index(int) .", textPred, popPred)) + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <300001> <%[1]s> "fox dog" . + <300002> <%[1]s> "fox cat bird" . + <300003> <%[1]s> "dog cat" . + <300004> <%[1]s> "fox fox" . + <300006> <%[1]s> "fox fox fox" . + <300001> <%[2]s> "1" . + <300002> <%[2]s> "2" . + <300003> <%[2]s> "3" . + <300004> <%[2]s> "-5" . + <300006> <%[2]s> "-7" . + `, textPred, popPred))) +} + +// KNOWN-FAILING (Q1): bm25 + @filter + first returns the N lowest-uid matches instead of the N highest-scored. +// +// With a non-uid @filter present, First is not pushed to the worker and the +// coordinator guard that skips applyPagination for bm25 roots only fires when there +// are no filters, so first:N is applied by x.PageRange over the uid-sorted uidMatrix +// — silently flipping the page from score order to uid order. +func TestBM25FilterThenFirstPreservesScoreOrder(t *testing.T) { + textPred, popPred := "bm25_q1_text", "bm25_q1_pop" + setupBM25FilterCorpus(t, textPred, popPred) + + js := processQueryNoErr(t, fmt.Sprintf(` + { + me(func: bm25(%s, "fox"), first: 2) @filter(ge(%s, -100)) { + uid + } + }`, textPred, popPred)) + uids, _ := parseUIDScoreRows(t, js) + require.ElementsMatch(t, []uint64{300006, 300004}, uids, + "first:2 with an all-passing filter must return the top-2 by score "+ + "(300006 'fox fox fox' and 300004 'fox fox'), not the 2 lowest-uid matches") +} + +// KNOWN-FAILING (Q1): bm25 + @filter + first/offset windows the uid-sorted matrix, not the score ranking. +// +// Companion to TestBM25FilterThenFirstPreservesScoreOrder: the offset:1 window must +// slice the score-descending ranking (score ranks 2..3), not the uid-ascending one. +func TestBM25FilterThenFirstOffsetWindowIsScoreRanked(t *testing.T) { + textPred, popPred := "bm25_q1o_text", "bm25_q1o_pop" + setupBM25FilterCorpus(t, textPred, popPred) + + js := processQueryNoErr(t, fmt.Sprintf(` + { + me(func: bm25(%s, "fox"), first: 2, offset: 1) @filter(ge(%s, -100)) { + uid + } + }`, textPred, popPred)) + uids, _ := parseUIDScoreRows(t, js) + require.ElementsMatch(t, []uint64{300004, 300001}, uids, + "first:2 offset:1 with an all-passing filter must return score ranks 2..3 "+ + "(300004 'fox fox' and 300001 'fox dog')") +} + +// TestBM25FilterDroppingTopKeepsDeeperMatches pins that the worker stays UNWINDOWED +// when a filter is present: a filter that removes the entire unfiltered top of the +// ranking must still surface the deeper matches. If a fix for Q1 ever pushes First +// down alongside a filter, the worker would retain only the (later filtered-out) +// top-2 and this page would come back empty. Passes today. +func TestBM25FilterDroppingTopKeepsDeeperMatches(t *testing.T) { + textPred, popPred := "bm25_q1d_text", "bm25_q1d_pop" + setupBM25FilterCorpus(t, textPred, popPred) + + // ge(pop, 0) drops 300006 and 300004 — the two top-scored fox docs. + js := processQueryNoErr(t, fmt.Sprintf(` + { + me(func: bm25(%s, "fox"), first: 2) @filter(ge(%s, 0)) { + uid + } + }`, textPred, popPred)) + uids, _ := parseUIDScoreRows(t, js) + require.ElementsMatch(t, []uint64{300001, 300002}, uids, + "when the filter drops the entire unfiltered top-2, the deeper matches "+ + "300001 and 300002 must still be returned (worker must not pre-window)") +} + +// TestBM25TieStormPaginationPartition walks a page window over a corpus of +// byte-identical documents (every score is an exact tie) and asserts the pages form a +// clean partition of the deterministic total order (score desc, uid asc — which under +// a full tie is just uid asc), with bit-identical scores everywhere and identical +// output on repeated walks. +// +// Q3 pin: the tie semantics (higher uid is the eviction victim, ties resolve uid-asc) +// is the current documented design of topKHeap; this test pins that behavior +// end-to-end. Any page overlap/gap or cross-walk score bit-difference decides Q3 to +// VULNERABLE (float-addition grouping differing between the per-page WAND paths). +func TestBM25TieStormPaginationPartition(t *testing.T) { + ensureBM25UidLease(t) + pred := "bm25_tiestorm" + t.Cleanup(func() { dropPredicate(pred) }) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + const ( + docBase = 310001 + numDocs = 300 + pageSize = 20 + walkRepeats = 30 + ) + + // 300 byte-identical docs matching a 3-term query. + for start := 0; start < numDocs; start += 100 { + var b strings.Builder + for i := start; i < start+100; i++ { + fmt.Fprintf(&b, "<%d> <%s> \"zephyrite quillon marblewood\" .\n", docBase+i, pred) + } + require.NoError(t, addTriplesToCluster(b.String())) + } + + pageQuery := fmt.Sprintf(` + { + score as var(func: bm25(%s, "zephyrite quillon marblewood"), first: %d, offset: %%d) + me(func: uid(score)) { + uid + val(score) + } + }`, pred, pageSize) + + walk := func() ([]uint64, map[uint64]float64) { + var all []uint64 + scores := make(map[uint64]float64, numDocs) + for offset := 0; offset < numDocs; offset += pageSize { + uids, pageScores := parseUIDScoreRows(t, + processQueryNoErr(t, fmt.Sprintf(pageQuery, offset))) + require.Len(t, uids, pageSize, "page at offset %d must be full", offset) + all = append(all, uids...) + for uid, s := range pageScores { + scores[uid] = s + } + } + return all, scores + } + + firstWalk, firstScores := walk() + + // The pages must partition the corpus exactly: under a full score tie the total + // order is uid ascending, so the concatenated pages are the uid-sorted corpus + // with no duplicate and no missing uid. + expected := make([]uint64, numDocs) + for i := range expected { + expected[i] = uint64(docBase + i) + } + require.Equal(t, expected, firstWalk, + "tie-storm pages must partition the corpus in (score desc, uid asc) order "+ + "with no duplicated or missing uid") + + // Byte-identical documents must produce bit-identical scores on every page. + plateau := firstScores[docBase] + for uid, s := range firstScores { + require.Equal(t, plateau, s, + "uid %d: byte-identical docs must have bit-identical scores across pages", uid) + } + + // Repeated walks must be byte-for-byte deterministic. + for r := 1; r < walkRepeats; r++ { + uids, scores := walk() + require.Equal(t, firstWalk, uids, "walk %d: page partition must be deterministic", r) + require.Equal(t, firstScores, scores, "walk %d: scores must be bit-identical", r) + } +} + +// TestBM25TieStormDeepOffset asserts a deep page (first:10, offset:10000) over an +// ~11k all-tie corpus returns exactly the true rank-10000..10009 documents. The +// worker sizes its top-k heap at first+offset for windowed queries; under a full tie +// the true ranks are uid-ascending, so the page is exactly uids at positions +// 10000..10009. Q3 pin — passes if the tie eviction boundary is exact at depth. +func TestBM25TieStormDeepOffset(t *testing.T) { + ensureBM25UidLease(t) + pred := "bm25_deepoffset" + t.Cleanup(func() { dropPredicate(pred) }) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + const ( + docBase = 320001 + numDocs = 11000 + ) + for start := 0; start < numDocs; start += 1000 { + var b strings.Builder + for i := start; i < start+1000; i++ { + fmt.Fprintf(&b, "<%d> <%s> \"abyssalite trenchwork lumenstone\" .\n", docBase+i, pred) + } + require.NoError(t, addTriplesToCluster(b.String())) + } + + js := processQueryNoErr(t, fmt.Sprintf(` + { + score as var(func: bm25(%s, "abyssalite trenchwork lumenstone"), first: 10, offset: 10000) + me(func: uid(score)) { + uid + val(score) + } + }`, pred)) + uids, _ := parseUIDScoreRows(t, js) + + expected := make([]uint64, 10) + for i := range expected { + expected[i] = uint64(docBase + 10000 + i) + } + require.Equal(t, expected, uids, + "first:10 offset:10000 must return exactly the true rank-10000..10009 uids") +} + +// TestBM25StatsIntegrityUnderResetChurn pins that idempotent re-SETs and value +// updates keep corpus stats exact: every SET is paired with a DEL of the found old +// value, so churn nets zero drift. Uses small single-doc transactions (isolating this +// from any intra-proposal batching race) on two uids congruent mod 32, so both docs +// share one stats bucket. After each phase the observed bm25 score must equal the +// closed-form value computed from the exact live-corpus (docCount, totalTerms). +// Q13 regression pin — passes today; any drift indicates a new unpaired-DEL path. +func TestBM25StatsIntegrityUnderResetChurn(t *testing.T) { + ensureBM25UidLease(t) + pred := "bm25_churn" + t.Cleanup(func() { dropPredicate(pred) }) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + // 340010 % 32 == 340042 % 32 == 10: both docs live in the same stats bucket. + const uidA, uidB = 340010, 340042 + + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <%d> <%s> "kraken kraken kraken" . + `, uidA, pred))) + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <%d> <%s> "kraken squid" . + `, uidB, pred))) + + // Closed-form BM25 with default k=1.2, b=0.75 for the single query term + // "kraken", which both docs always contain (df=2). + closedForm := func(tf, dl, docCount, totalTerms float64) float64 { + const k, b, df = 1.2, 0.75, 2.0 + idf := math.Log1p((docCount - df + 0.5) / (df + 0.5)) + avgDL := totalTerms / docCount + return idf * (k + 1) * tf / (k*(1-b+b*dl/avgDL) + tf) + } + + scoreQuery := fmt.Sprintf(` + { + score as var(func: bm25(%s, "kraken")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }`, pred) + + assertScores := func(phase string, expA, expB float64) { + t.Helper() + _, scores := parseUIDScoreRows(t, processQueryNoErr(t, scoreQuery)) + require.Len(t, scores, 2, "%s: both docs must match 'kraken'", phase) + require.InEpsilon(t, expA, scores[uidA], 1e-9, + "%s: doc A score must match closed-form from exact (docCount, totalTerms)", phase) + require.InEpsilon(t, expB, scores[uidB], 1e-9, + "%s: doc B score must match closed-form from exact (docCount, totalTerms)", phase) + } + + // Baseline: docCount=2, totalTerms=5 (dl 3 + 2). + assertScores("baseline", closedForm(3, 3, 2, 5), closedForm(1, 2, 2, 5)) + + // Phase 1: re-SET the identical triple 10x in single-doc transactions. + for i := 0; i < 10; i++ { + require.NoError(t, addTriplesToCluster(fmt.Sprintf( + `<%d> <%s> "kraken kraken kraken" .`, uidA, pred))) + } + assertScores("after idempotent re-SETs", closedForm(3, 3, 2, 5), closedForm(1, 2, 2, 5)) + + // Phase 2: update doc B 10x, alternating lengths, ending on "kraken squid ink" + // (dl=3): docCount=2, totalTerms=6. + for i := 0; i < 10; i++ { + content := "kraken squid" + if i%2 == 1 { + content = "kraken squid ink" + } + require.NoError(t, addTriplesToCluster(fmt.Sprintf( + `<%d> <%s> %q .`, uidB, pred, content))) + } + assertScores("after update churn", closedForm(3, 3, 2, 6), closedForm(1, 3, 2, 6)) +} + +// Q27: bm25 has no @lang semantics — lang-tagged values are indexed under +// lang-qualified keys that bm25() cannot query (empirically: the tagged value is +// stored and readable via pred@de, but invisible to bm25() and its stats are +// asymmetric on 'S P *' delete, which re-tokenizes with the default analyzer). The +// schema combination is therefore rejected outright (schema/parse.go), matching the +// existing @noconflict and list-predicate rejections. This test pins the rejection; +// if bm25 ever grows real @lang support, replace it with the symmetric index/delete +// battery (see the Q27 verdict in the testing report). +func TestBM25LangSchemaRejected(t *testing.T) { + pred := "bm25_langdel" + t.Cleanup(func() { dropPredicate(pred) }) + err := client.Alter(context.Background(), + &api.Operation{Schema: fmt.Sprintf("%s: string @lang @index(bm25) .", pred)}) + require.Error(t, err, "@lang + @index(bm25) must be rejected at schema time") + require.Contains(t, err.Error(), "bm25 does not support language-qualified values") +} + +// TestBM25EmojiOnlyValueStatsSymmetry pins the symmetric-safe zero-token path (Q27 +// companion): a value that tokenizes to zero terms contributes nothing to corpus +// stats on SET and nothing on DEL — the docLen==0 skip applies to both directions — +// so every other document's score is untouched by its whole lifecycle. Passes today. +func TestBM25EmojiOnlyValueStatsSymmetry(t *testing.T) { + ensureBM25UidLease(t) + pred := "bm25_emoji" + t.Cleanup(func() { dropPredicate(pred) }) + setSchema(fmt.Sprintf("%s: string @index(bm25) .", pred)) + + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <350011> <%[1]s> "gondola gondola" . + <350012> <%[1]s> "gondola vellum" . + `, pred))) + + scoreQuery := fmt.Sprintf(` + { + score as var(func: bm25(%s, "gondola")) + me(func: uid(score), orderdesc: val(score)) { + uid + val(score) + } + }`, pred) + _, baseline := parseUIDScoreRows(t, processQueryNoErr(t, scoreQuery)) + baseScore, ok := baseline[350011] + require.True(t, ok, "control doc 350011 must be in the baseline results") + + // An emoji-only value tokenizes to zero terms and must be a stats no-op. + require.NoError(t, addTriplesToCluster(fmt.Sprintf( + `<350013> <%s> "🦊🐈🦄" .`, pred))) + _, afterSet := parseUIDScoreRows(t, processQueryNoErr(t, scoreQuery)) + require.Equal(t, baseScore, afterSet[350011], + "an emoji-only (zero-token) SET must not change corpus stats") + + deleteTriplesInCluster(fmt.Sprintf(`<350013> <%s> * .`, pred)) + _, afterDel := parseUIDScoreRows(t, processQueryNoErr(t, scoreQuery)) + require.Equal(t, baseScore, afterDel[350011], + "deleting an emoji-only (zero-token) value must not change corpus stats") +} diff --git a/query/query_hybrid_test.go b/query/query_hybrid_test.go new file mode 100644 index 00000000000..2f0e4ed7ad2 --- /dev/null +++ b/query/query_hybrid_test.go @@ -0,0 +1,1105 @@ +//go:build integration || cloud + +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +//nolint:lll +package query + +import ( + "context" + "encoding/json" + "fmt" + "math" + "sort" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +// hybridResult unmarshals a {uid, val(f)} result list ordered by fused score. +type hybridRow struct { + UID string `json:"uid"` + Score float64 `json:"val(f)"` +} + +func fuseRows(t *testing.T, js string) []hybridRow { + t.Helper() + var resp struct { + Data struct { + Me []hybridRow `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + return resp.Data.Me +} + +func fuseUIDSet(rows []hybridRow) map[string]bool { + set := make(map[string]bool, len(rows)) + for _, r := range rows { + set[r.UID] = true + } + return set +} + +// --- similar_to score surfacing (prerequisite for vector fusion) ------------- + +func TestSimilarToScoreVariable(t *testing.T) { + // similar_to bound to a value variable surfaces a higher-is-better similarity + // score. The query vector equals doc 503's embedding, so 503 is the closest + // (euclidean distance 0 -> score 1.0) and must rank first. + query := ` + { + s as var(func: similar_to(description_vec, 7, "[3.0, 0.0, 0.0, 0.0]")) + me(func: uid(s), orderdesc: val(s)) { + uid + val(s) + } + }` + js := processQueryNoErr(t, query) + var resp struct { + Data struct { + Me []struct { + UID string `json:"uid"` + Score float64 `json:"val(s)"` + } `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + require.NotEmpty(t, resp.Data.Me) + + uid503 := uidHex(t, 503) + require.Equal(t, uid503, resp.Data.Me[0].UID, "closest vector (503) should rank first") + require.InDelta(t, 1.0, resp.Data.Me[0].Score, 1e-6, "exact match should score 1/(1+0)=1.0") + + // Scores must be in descending order. + for i := 1; i < len(resp.Data.Me); i++ { + require.GreaterOrEqual(t, resp.Data.Me[i-1].Score, resp.Data.Me[i].Score) + } +} + +// --- fuse() over BM25 channels ---------------------------------------------- + +func TestFuseRRFTwoBM25Channels(t *testing.T) { + // Fuse two BM25 channels: "fox" (matches 501,502,503,506,507) and "dog" + // (matches 501,502,504,505). The union must include dog-only docs (504,505) + // that the fox channel never returns — proving outer-join (not intersection). + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + dog as var(func: bm25(description_bm25, "dog")) + f as var(func: fuse(fox, dog, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { + uid + val(f) + } + }` + rows := fuseRows(t, processQueryNoErr(t, query)) + require.NotEmpty(t, rows) + set := fuseUIDSet(rows) + + // Union contains both fox-only (503) and dog-only (504, 505) documents. + require.True(t, set[uidHex(t, 503)], "fox-only doc 503 must be present") + require.True(t, set[uidHex(t, 504)], "dog-only doc 504 must be present (union, not intersection)") + require.True(t, set[uidHex(t, 505)], "dog-only doc 505 must be present") + + // Doc 501/502 appear in BOTH channels, so their fused RRF score should exceed a + // document that appears in only one channel at the same rank. + scores := make(map[string]float64) + for _, r := range rows { + scores[r.UID] = r.Score + } + require.Greater(t, scores[uidHex(t, 501)], 0.0) + + // Scores descending. + for i := 1; i < len(rows); i++ { + require.GreaterOrEqual(t, rows[i-1].Score, rows[i].Score) + } +} + +func TestFuseLinearWeights(t *testing.T) { + // Linear fusion with weights. Heavily weight the "dog" channel; a dog-only doc + // should then be present and outscore where appropriate. + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + dog as var(func: bm25(description_bm25, "dog")) + f as var(func: fuse(fox, dog, method: "linear", weights: "0.1,0.9", normalize: "max")) + me(func: uid(f), orderdesc: val(f)) { + uid + val(f) + } + }` + rows := fuseRows(t, processQueryNoErr(t, query)) + require.NotEmpty(t, rows) + set := fuseUIDSet(rows) + require.True(t, set[uidHex(t, 504)], "dog-only doc must appear with linear fusion") + for i := 1; i < len(rows); i++ { + require.GreaterOrEqual(t, rows[i-1].Score, rows[i].Score) + } +} + +func TestFuseSingleChannelPassthrough(t *testing.T) { + // A single-channel fuse should preserve that channel's ranking. "fox fox fox" + // (503) is the top BM25 result for "fox". + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(fox, method: "rrf")) + me(func: uid(f), orderdesc: val(f), first: 1) { + uid + val(f) + } + }` + rows := fuseRows(t, processQueryNoErr(t, query)) + require.Len(t, rows, 1) + require.Equal(t, uidHex(t, 503), rows[0].UID) +} + +func TestFusePagination(t *testing.T) { + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + dog as var(func: bm25(description_bm25, "dog")) + f as var(func: fuse(fox, dog, method: "rrf")) + me(func: uid(f), orderdesc: val(f), first: 2, offset: 1) { + uid + val(f) + } + }` + rows := fuseRows(t, processQueryNoErr(t, query)) + require.Len(t, rows, 2, "first:2 offset:1 should return exactly 2 rows") +} + +// --- fuse() hybrid: BM25 + vector ------------------------------------------- + +func TestFuseHybridBM25AndVector(t *testing.T) { + // The headline use case: fuse BM25 text relevance with vector similarity. + // Doc 503 ("fox fox fox", embedding [3,0,0,0]) is rank-1 in BOTH the "fox" BM25 + // channel and the [3,0,0,0] vector channel, so it must be the top fused result. + // The vector channel (k=7) returns all docs, so dog-only docs (504,505) enter + // the union even though the BM25 "fox" channel never returns them. + query := ` + { + txt as var(func: bm25(description_bm25, "fox")) + vec as var(func: similar_to(description_vec, 7, "[3.0, 0.0, 0.0, 0.0]")) + f as var(func: fuse(txt, vec, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { + uid + val(f) + } + }` + rows := fuseRows(t, processQueryNoErr(t, query)) + require.NotEmpty(t, rows) + require.Equal(t, uidHex(t, 503), rows[0].UID, "503 is rank-1 in both channels -> top fused") + + set := fuseUIDSet(rows) + require.True(t, set[uidHex(t, 504)], "vector-only doc 504 must enter the union") + require.True(t, set[uidHex(t, 505)], "vector-only doc 505 must enter the union") + + for i := 1; i < len(rows); i++ { + require.GreaterOrEqual(t, rows[i-1].Score, rows[i].Score) + } +} + +// --- hybrid() sugar ---------------------------------------------------------- + +func TestHybridSugarEquivalentToFuse(t *testing.T) { + // hybrid() must produce the same top result as the explicit fuse() form. + hybridQ := ` + { + f as var(func: hybrid(description_bm25, "fox", description_vec, "[3.0, 0.0, 0.0, 0.0]", topk: 7, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { + uid + val(f) + } + }` + explicitQ := ` + { + txt as var(func: bm25(description_bm25, "fox"), first: 7) + vec as var(func: similar_to(description_vec, 7, "[3.0, 0.0, 0.0, 0.0]")) + f as var(func: fuse(txt, vec, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { + uid + val(f) + } + }` + hybridRows := fuseRows(t, processQueryNoErr(t, hybridQ)) + explicitRows := fuseRows(t, processQueryNoErr(t, explicitQ)) + + require.NotEmpty(t, hybridRows) + require.Equal(t, len(explicitRows), len(hybridRows), "hybrid and fuse must return the same set") + require.Equal(t, explicitRows[0].UID, hybridRows[0].UID, "same top result") + // Fused scores should match position-for-position. + for i := range explicitRows { + require.Equal(t, explicitRows[i].UID, hybridRows[i].UID, "row %d uid mismatch", i) + require.InDelta(t, explicitRows[i].Score, hybridRows[i].Score, 1e-9, "row %d score mismatch", i) + } +} + +// --- error handling ---------------------------------------------------------- + +func TestFuseUnknownMethod(t *testing.T) { + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(fox, method: "bogus")) + me(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := processQuery(context.Background(), t, query) + require.Error(t, err) + require.Contains(t, err.Error(), "method") +} + +func TestFuseWeightsCountMismatch(t *testing.T) { + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + dog as var(func: bm25(description_bm25, "dog")) + f as var(func: fuse(fox, dog, method: "linear", weights: "0.5")) + me(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := processQuery(context.Background(), t, query) + require.Error(t, err) + require.Contains(t, err.Error(), "weights") +} + +func TestFuseBadK(t *testing.T) { + query := ` + { + fox as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(fox, k: "-5")) + me(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := processQuery(context.Background(), t, query) + require.Error(t, err) + require.Contains(t, err.Error(), "k must be") +} + +// --- generic block helpers ---------------------------------------------------- + +// valRow is a {uid, val()} row from an arbitrary result block. +type valRow struct { + UID string + Score float64 +} + +// blockValRows extracts the ordered {uid, val(varName)} rows of the named block. +// A missing val() key yields Score 0 (callers assert presence separately). +func blockValRows(t *testing.T, js, block, varName string) []valRow { + t.Helper() + var resp struct { + Data map[string]json.RawMessage `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp), "response must be valid JSON: %s", js) + raw, ok := resp.Data[block] + require.Truef(t, ok, "block %q missing from response: %s", block, js) + var items []map[string]interface{} + require.NoError(t, json.Unmarshal(raw, &items)) + key := "val(" + varName + ")" + rows := make([]valRow, 0, len(items)) + for _, it := range items { + uid, _ := it["uid"].(string) + score, _ := it[key].(float64) + rows = append(rows, valRow{UID: uid, Score: score}) + } + return rows +} + +// blockScoreMap is blockValRows collapsed into a uid->score map. +func blockScoreMap(t *testing.T, js, block, varName string) map[string]float64 { + t.Helper() + rows := blockValRows(t, js, block, varName) + m := make(map[string]float64, len(rows)) + for _, r := range rows { + m[r.UID] = r.Score + } + return m +} + +// blockUids extracts the ordered uid list of the named block. +func blockUids(t *testing.T, js, block string) []string { + t.Helper() + rows := blockValRows(t, js, block, "") + uids := make([]string, 0, len(rows)) + for _, r := range rows { + uids = append(uids, r.UID) + } + return uids +} + +// uidNum parses a hex uid string ("0x1f5") into its numeric value. +func uidNum(uid string) uint64 { + v, err := strconv.ParseUint(strings.TrimPrefix(uid, "0x"), 16, 64) + if err != nil { + panic(fmt.Sprintf("bad uid %q: %v", uid, err)) + } + return v +} + +// rrfExpected re-implements the fusion contract of query/fuse.go (fuseRRF + +// channelRanks: rank by score desc, tie-break uid asc; contribution +// weight * 1/(k+rank)) so e2e results can be checked against an independent oracle. +func rrfExpected(channels []map[string]float64, weights []float64, k float64) map[string]float64 { + fused := make(map[string]float64) + for ci, ch := range channels { + uids := make([]string, 0, len(ch)) + for uid := range ch { + uids = append(uids, uid) + } + sort.Slice(uids, func(i, j int) bool { + si, sj := ch[uids[i]], ch[uids[j]] + if si != sj { + return si > sj + } + return uidNum(uids[i]) < uidNum(uids[j]) + }) + w := 1.0 + if weights != nil { + w = weights[ci] + } + for r, uid := range uids { + fused[uid] += w * (1.0 / (k + float64(r+1))) + } + } + return fused +} + +// --- Q2: ranker score binding under filter + pagination + traversal ----------- + +// TestRankerScoreBindingUnderFilterPaginationTraversal is a regression pin for the +// uid-keyed rankerScores snapshot (commits 250b79bec/98ae4335e): per-uid scores +// bound by bm25/similar_to must stay attached to the right uid after a var-level +// @filter shrinks the matched set, after orderdesc+first+offset pagination, and +// when the same doc is reached through multiple traversal paths. +func TestRankerScoreBindingUnderFilterPaginationTraversal(t *testing.T) { + const ( + textPred = "hyb_bind_txt" + linkPred = "hyb_bind_link" + vecPred = "hyb_bind_vec" + ) + setSchema(fmt.Sprintf(` + %s: string @index(bm25) . + %s: [uid] . + %s: float32vector @index(hnsw(metric:"euclidean")) . + `, textPred, linkPred, vecPred)) + t.Cleanup(func() { + dropPredicate(textPred) + dropPredicate(linkPred) + dropPredicate(vecPred) + }) + + // 8 text docs, tf("target")=1 each with strictly increasing doclen, so the 8 + // BM25 scores are strictly distinct and rank 7101 > 7102 > ... > 7108. + // 8 vector docs at distance i from the origin, so similarity 1/(1+i) is + // strictly distinct and rank 7121 > 7122 > ... > 7128. + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <7101> <%[1]s> "target" . + <7102> <%[1]s> "target alpha" . + <7103> <%[1]s> "target alpha bravo" . + <7104> <%[1]s> "target alpha bravo charlie" . + <7105> <%[1]s> "target alpha bravo charlie delta" . + <7106> <%[1]s> "target alpha bravo charlie delta echo" . + <7107> <%[1]s> "target alpha bravo charlie delta echo foxtrot" . + <7108> <%[1]s> "target alpha bravo charlie delta echo foxtrot golf" . + <7111> <%[2]s> <7103> . + <7111> <%[2]s> <7105> . + <7112> <%[2]s> <7103> . + <7112> <%[2]s> <7105> . + <7121> <%[3]s> "[1.0, 0.0]" . + <7122> <%[3]s> "[2.0, 0.0]" . + <7123> <%[3]s> "[3.0, 0.0]" . + <7124> <%[3]s> "[4.0, 0.0]" . + <7125> <%[3]s> "[5.0, 0.0]" . + <7126> <%[3]s> "[6.0, 0.0]" . + <7127> <%[3]s> "[7.0, 0.0]" . + <7128> <%[3]s> "[8.0, 0.0]" . + `, textPred, linkPred, vecPred))) + + t.Run("BM25", func(t *testing.T) { + // Unfiltered control: the per-uid score oracle. + oracleJs := processQueryNoErr(t, fmt.Sprintf(` + { + s as var(func: bm25(%s, "target")) + all(func: uid(s), orderdesc: val(s)) { uid val(s) } + }`, textPred)) + oracleRows := blockValRows(t, oracleJs, "all", "s") + require.Len(t, oracleRows, 8) + for i := 1; i < len(oracleRows); i++ { + require.Greater(t, oracleRows[i-1].Score, oracleRows[i].Score, + "scores must be strictly distinct for an unambiguous oracle") + } + require.Equal(t, uidHex(t, 7101), oracleRows[0].UID, "shortest doc must rank first") + oracle := make(map[string]float64, 8) + for _, r := range oracleRows { + oracle[r.UID] = r.Score + } + + // Filter to alternating ranks, then paginate and traverse. Every emitted + // val(s) must equal the oracle score of the SAME uid (uid-keyed comparison). + js := processQueryNoErr(t, fmt.Sprintf(` + { + s as var(func: bm25(%[1]s, "target")) @filter(uid(7101, 7103, 7105, 7107)) + page(func: uid(s), orderdesc: val(s), first: 2, offset: 1) { uid val(s) } + parents(func: uid(7111, 7112)) { uid %[2]s { uid val(s) } } + }`, textPred, linkPred)) + + page := blockValRows(t, js, "page", "s") + require.Len(t, page, 2) + // Filtered ranking is 7101 > 7103 > 7105 > 7107; offset:1 first:2 -> 7103, 7105. + require.Equal(t, uidHex(t, 7103), page[0].UID) + require.Equal(t, uidHex(t, 7105), page[1].UID) + for _, r := range page { + require.Equal(t, oracle[r.UID], r.Score, + "paginated val(s) for %s must equal the unfiltered control score", r.UID) + } + + // Each doc is reached via TWO parent paths; both must resolve the same score. + var presp struct { + Data struct { + Parents []struct { + UID string `json:"uid"` + Link []struct { + UID string `json:"uid"` + Score float64 `json:"val(s)"` + } `json:"hyb_bind_link"` + } `json:"parents"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &presp)) + require.Len(t, presp.Data.Parents, 2) + for _, p := range presp.Data.Parents { + require.Len(t, p.Link, 2, "parent %s must reach both docs", p.UID) + for _, ch := range p.Link { + require.Equal(t, oracle[ch.UID], ch.Score, + "traversal val(s) for %s via parent %s must equal the control score", ch.UID, p.UID) + } + } + }) + + t.Run("SimilarTo", func(t *testing.T) { + oracleJs := processQueryNoErr(t, fmt.Sprintf(` + { + s as var(func: similar_to(%s, 10, "[0.0, 0.0]")) + all(func: uid(s), orderdesc: val(s)) { uid val(s) } + }`, vecPred)) + oracleRows := blockValRows(t, oracleJs, "all", "s") + require.Len(t, oracleRows, 8) + oracle := make(map[string]float64, 8) + for _, r := range oracleRows { + oracle[r.UID] = r.Score + } + // Metric-math cross-check: doc 712i sits at distance i, score 1/(1+i). + for i := 1; i <= 8; i++ { + require.InDelta(t, 1.0/(1.0+float64(i)), oracle[uidHex(t, 7120+i)], 1e-6) + } + + js := processQueryNoErr(t, fmt.Sprintf(` + { + s as var(func: similar_to(%s, 10, "[0.0, 0.0]")) @filter(uid(7121, 7123, 7125, 7127)) + page(func: uid(s), orderdesc: val(s), first: 2, offset: 1) { uid val(s) } + }`, vecPred)) + page := blockValRows(t, js, "page", "s") + require.Len(t, page, 2) + // Filtered ranking is 7121 > 7123 > 7125 > 7127; offset:1 first:2 -> 7123, 7125. + require.Equal(t, uidHex(t, 7123), page[0].UID) + require.Equal(t, uidHex(t, 7125), page[1].UID) + for _, r := range page { + require.Equal(t, oracle[r.UID], r.Score, + "paginated val(s) for %s must equal the unfiltered control score", r.UID) + } + }) +} + +// --- Q4: similar_to score orientation matches the metric math ----------------- + +// TestSimilarToScoreValueMatchesMetricMath pins the score orientation contract of +// tok/hnsw similarityScore: cosine/dot surface the raw similarity (higher-better, +// as-is) and euclidean surfaces 1/(1+d) with d the TRUE (non-squared) L2 distance. +// The vectors are chosen so the three metrics rank the corpus differently, so a +// future sign flip, double transform, or squared-vs-true-distance regression +// cannot slip through. A single-channel fuse() consuming the same variable must +// leave the channel's surfaced scores untouched. +func TestSimilarToScoreValueMatchesMetricMath(t *testing.T) { + preds := map[string]string{ + "euclidean": "hyb_metric_euc", + "cosine": "hyb_metric_cos", + "dotproduct": "hyb_metric_dot", + } + setSchema(fmt.Sprintf(` + %s: float32vector @index(hnsw(metric:"euclidean")) . + %s: float32vector @index(hnsw(metric:"cosine")) . + %s: float32vector @index(hnsw(metric:"dotproduct")) . + `, preds["euclidean"], preds["cosine"], preds["dotproduct"])) + t.Cleanup(func() { + for _, p := range preds { + dropPredicate(p) + } + }) + + // Same four vectors on all three predicates; query vector is [1, 0]. + // 7141: [1, 0] 7142: [0.6, 0.8] 7143: [4, 0] 7144: [0, 1] + var triples strings.Builder + for _, p := range preds { + fmt.Fprintf(&triples, ` + <7141> <%[1]s> "[1.0, 0.0]" . + <7142> <%[1]s> "[0.6, 0.8]" . + <7143> <%[1]s> "[4.0, 0.0]" . + <7144> <%[1]s> "[0.0, 1.0]" .`, p) + } + require.NoError(t, addTriplesToCluster(triples.String())) + + expected := map[string]map[int]float64{ + // 1/(1+d) with d the true L2 distance to [1,0]. + "euclidean": { + 7141: 1.0, + 7142: 1.0 / (1.0 + math.Sqrt(0.8)), + 7143: 1.0 / (1.0 + 3.0), + 7144: 1.0 / (1.0 + math.Sqrt2), + }, + // Raw cosine similarity, as-is. + "cosine": {7141: 1.0, 7142: 0.6, 7143: 1.0, 7144: 0.0}, + // Raw dot product, as-is (unbounded, higher-better). + "dotproduct": {7141: 1.0, 7142: 0.6, 7143: 4.0, 7144: 0.0}, + } + + scores := make(map[string]map[string]float64) + for metric, pred := range preds { + js := processQueryNoErr(t, fmt.Sprintf(` + { + s as var(func: similar_to(%s, 10, "[1.0, 0.0]")) + ranked(func: uid(s), orderdesc: val(s)) { uid val(s) } + f as var(func: fuse(s, method: "rrf")) + fused(func: uid(f), orderdesc: val(f)) { uid } + }`, pred)) + got := blockScoreMap(t, js, "ranked", "s") + require.Len(t, got, 4, "%s: all four vectors must be returned", metric) + for decUID, want := range expected[metric] { + require.InDelta(t, want, got[uidHex(t, decUID)], 1e-6, + "%s: score for uid %d must match the hand-computed metric value", metric, decUID) + } + // fuse() consuming the channel must not perturb it: the fused uid set is the + // channel's uid set and the surfaced channel scores stay hand-computed-exact. + require.ElementsMatch(t, + []string{uidHex(t, 7141), uidHex(t, 7142), uidHex(t, 7143), uidHex(t, 7144)}, + blockUids(t, js, "fused"), "%s: fused uid set must equal the channel's", metric) + scores[metric] = got + } + + // The three metrics must genuinely disagree on the ranking (otherwise this test + // couldn't catch an orientation mix-up): euclidean ranks 7143 last-but-one, + // dotproduct ranks it first. + require.Greater(t, scores["euclidean"][uidHex(t, 7141)], scores["euclidean"][uidHex(t, 7143)]) + require.Greater(t, scores["dotproduct"][uidHex(t, 7143)], scores["dotproduct"][uidHex(t, 7141)]) +} + +// --- Q5: NaN/Inf query vectors ------------------------------------------------ + +// KNOWN-FAILING (Q5): a NaN query vector parses silently, poisons every channel score with NaN, and fuse drops the whole vector channel — the hybrid result is byte-identical to a bm25-only control instead of erroring. +// Correct behavior asserted here: a non-finite query vector either fails with a +// clean query error, or produces valid JSON with finite scores, deterministic +// ordering, and a vector channel that actually participates in fusion. +func TestSimilarToNaNInfQueryVector(t *testing.T) { + // bm25-only control for the channel-participation check. + controlJs := processQueryNoErr(t, ` + { + txt as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(txt, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`) + controlRows := blockValRows(t, controlJs, "me", "f") + require.NotEmpty(t, controlRows) + + for _, tc := range []struct{ name, vec string }{ + {"NaN", "[NaN, 0.0, 0.0, 0.0]"}, + {"PlusInf", "[Inf, 0.0, 0.0, 0.0]"}, + } { + t.Run(tc.name, func(t *testing.T) { + // Standalone similar_to: either a clean error, or valid JSON with finite scores. + standalone := fmt.Sprintf(` + { + s as var(func: similar_to(description_vec, 7, "%s")) + me(func: uid(s), orderdesc: val(s)) { uid val(s) } + }`, tc.vec) + js, err := processQuery(context.Background(), t, standalone) + if err == nil { + rows := blockValRows(t, js, "me", "s") // fails on a bare NaN token in the JSON + for _, r := range rows { + require.False(t, math.IsNaN(r.Score) || math.IsInf(r.Score, 0), + "surfaced score for %s must be finite", r.UID) + } + } + + // Fused: the vector channel must not be silently dropped. If the query is + // accepted, the fused ranking must differ from the bm25-only control + // (the vector channel contributes ranks even when degenerate) and must be + // deterministic across runs. + mixed := fmt.Sprintf(` + { + txt as var(func: bm25(description_bm25, "fox")) + vec as var(func: similar_to(description_vec, 7, "%s")) + f as var(func: fuse(txt, vec, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`, tc.vec) + jm, err := processQuery(context.Background(), t, mixed) + if err != nil { + return // clean rejection is correct behavior + } + rows1 := blockValRows(t, jm, "me", "f") + rows2 := blockValRows(t, processQueryNoErr(t, mixed), "me", "f") + require.Equal(t, rows1, rows2, "degenerate-vector fusion must be deterministic") + require.NotEqual(t, controlRows, rows1, + "fuse(txt, vec) with a %s query vector must not silently equal the bm25-only control "+ + "(the vector channel was dropped without an error)", tc.name) + }) + } +} + +// --- Q6: hybrid topk semantics ------------------------------------------------ + +// TestHybridTopKIsChannelDepthNotFusedCap pins the CURRENT topk semantics of +// hybrid(): topk bounds each channel's retrieval depth (bm25 first + similar_to +// neighbor count), NOT the fused output size — the fused union of two depth-N +// channels may contain up to 2N uids. AMBIGUOUS-SEMANTICS NOTE: the design doc +// (2026-06-03-hybrid-search-fusion-design.md) describes topk as a cap on emitted +// fused results, and two-stage retrieval inherently omits a doc ranked N+1 in both +// channels even when its exhaustive fused score would win. Whether topk should be +// forwarded as a post-fusion cap (fuse's own topk option) is an open decision; +// this test pins the current deterministic desugaring so a semantics change is a +// deliberate, visible act. +func TestHybridTopKIsChannelDepthNotFusedCap(t *testing.T) { + hybridQ := ` + { + f as var(func: hybrid(description_bm25, "fox", description_vec, "[3.0, 0.0, 0.0, 0.0]", topk: 2, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { uid val(f) } + }` + explicitQ := ` + { + txt as var(func: bm25(description_bm25, "fox"), first: 2) + vec as var(func: similar_to(description_vec, 2, "[3.0, 0.0, 0.0, 0.0]")) + f as var(func: fuse(txt, vec, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { uid val(f) } + }` + + hybridRows := blockValRows(t, processQueryNoErr(t, hybridQ), "me", "f") + explicitRows := blockValRows(t, processQueryNoErr(t, explicitQ), "me", "f") + + // Deterministic across runs. + hybridRows2 := blockValRows(t, processQueryNoErr(t, hybridQ), "me", "f") + require.Equal(t, hybridRows, hybridRows2, "hybrid(topk:2) must be deterministic") + + // hybrid desugars exactly to the two-stage explicit form. + require.Equal(t, len(explicitRows), len(hybridRows)) + for i := range explicitRows { + require.Equal(t, explicitRows[i].UID, hybridRows[i].UID, "row %d uid mismatch", i) + require.InDelta(t, explicitRows[i].Score, hybridRows[i].Score, 1e-9, "row %d score mismatch", i) + } + + // topk is NOT a post-fusion cap: the union of the depth-2 bm25 channel + // ({503, 501}) and the depth-2 vector channel ({503, 506}) has 3 uids. + require.Greater(t, len(hybridRows), 2, + "fused output must exceed topk (topk bounds per-channel retrieval, not fused emission)") + set := make(map[string]bool) + for _, r := range hybridRows { + set[r.UID] = true + } + require.True(t, set[uidHex(t, 503)], "503 is rank-1 in both depth-2 channels") +} + +// --- Q8: channel validity contract -------------------------------------------- + +// KNOWN-FAILING (Q8): fuse() silently treats a uid variable (which carries no scores) as a valid empty channel and returns the other channel's ranking with no error. +// Correct behavior asserted here: a non-ranker channel is rejected with an error +// naming the offending channel, mirroring scoresFromVar's "non-numeric scores" +// contract instead of silently no-op'ing the channel. +func TestFuseUidVarChannelRejected(t *testing.T) { + query := ` + { + u as var(func: has(description_bm25)) + fox as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(u, fox, method: "rrf", k: 60)) + me(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := processQuery(context.Background(), t, query) + require.Error(t, err, "a uid-var channel (no scores) must be rejected, not silently ignored") + require.Contains(t, err.Error(), "channel") +} + +// KNOWN-FAILING (Q8): a count(uid) variable used as a fuse channel injects the sentinel uid math.MaxUint64 (0xffffffffffffffff) into the fused output as a phantom result. +func TestFuseCountVarChannelNoSentinelUid(t *testing.T) { + query := ` + { + var(func: has(description_bm25)) { cnt as count(uid) } + fox as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(cnt, fox, method: "rrf", k: 60)) + me(func: uid(f)) { uid } + }` + js, err := processQuery(context.Background(), t, query) + if err != nil { + return // rejecting a count-var channel outright is also correct behavior + } + require.NotContains(t, js, "0xffffffffffffffff", + "the count(uid) sentinel uid must never surface as a fused result") +} + +// TestFuseEmptyChannelKeepsOtherChannels is a passing pin for the documented +// empty-channel semantics (query/fuse.go computeFuse): a ranker channel that ran +// but matched nothing contributes nothing and must not drop or reorder the other +// channels' results — the fused output equals the surviving channel's single- +// channel fusion exactly. +func TestFuseEmptyChannelKeepsOtherChannels(t *testing.T) { + js := processQueryNoErr(t, ` + { + dead as var(func: bm25(description_bm25, "zzzqqqxyzzy")) + fox as var(func: bm25(description_bm25, "fox")) + f as var(func: fuse(dead, fox, method: "rrf", k: 60)) + solo as var(func: fuse(fox, method: "rrf", k: 60)) + both(func: uid(f), orderdesc: val(f)) { uid val(f) } + alone(func: uid(solo), orderdesc: val(solo)) { uid val(solo) } + }`) + both := blockValRows(t, js, "both", "f") + alone := blockValRows(t, js, "alone", "solo") + require.NotEmpty(t, both) + require.Equal(t, alone, both, + "an empty channel must contribute nothing: fuse(dead, fox) == fuse(fox) exactly") +} + +// --- Q12: manual union + orderdesc over one channel's scores ------------------ + +// TestHybridUnionOrderByOneChannelDropsUnscored pins the CURRENT (inherited +// upstream) value-var ordering semantics for the manual-hybrid union pattern: +// `uid(b, v)` unions both channels, but `orderdesc: val(b)` sorts via +// sortAndPaginateUsingVar, which SILENTLY DROPS every uid that has no entry in +// b's score map — vector-only docs vanish from the ordered block even though the +// unordered union contains them. AMBIGUOUS-SEMANTICS NOTE: whether such uids +// should instead sort last (or error) is an open decision; this test pins the +// deterministic drop so any future semantics change is deliberate and visible. +func TestHybridUnionOrderByOneChannelDropsUnscored(t *testing.T) { + query := ` + { + b as var(func: bm25(description_bm25, "fox")) + v as var(func: similar_to(description_vec, 7, "[3.0, 0.0, 0.0, 0.0]")) + control(func: uid(b, v)) { uid } + ordered(func: uid(b, v), orderdesc: val(b)) { uid val(b) } + }` + js := processQueryNoErr(t, query) + + // The unordered union contains every doc from both channels, including the + // vector-only docs 504 and 505. + controlSet := make(map[string]bool) + for _, uid := range blockUids(t, js, "control") { + controlSet[uid] = true + } + for _, dec := range []int{501, 502, 503, 504, 505, 506, 507} { + require.True(t, controlSet[uidHex(t, dec)], "unordered union must contain %d", dec) + } + + // CURRENT behavior: orderdesc: val(b) drops the vector-only docs (absent from + // b's score map) instead of sorting them last. + ordered := blockValRows(t, js, "ordered", "b") + orderedSet := make(map[string]bool) + for _, r := range ordered { + orderedSet[r.UID] = true + } + require.Len(t, ordered, 5, "current semantics: only b-scored uids survive ordering") + for _, dec := range []int{501, 502, 503, 506, 507} { + require.True(t, orderedSet[uidHex(t, dec)], "bm25-scored doc %d must survive ordering", dec) + } + require.False(t, orderedSet[uidHex(t, 504)], "vector-only 504 is dropped (current semantics)") + require.False(t, orderedSet[uidHex(t, 505)], "vector-only 505 is dropped (current semantics)") + for i := 1; i < len(ordered); i++ { + require.GreaterOrEqual(t, ordered[i-1].Score, ordered[i].Score) + } + + // The drop must at least be deterministic. + js2 := processQueryNoErr(t, query) + require.Equal(t, ordered, blockValRows(t, js2, "ordered", "b"), + "the ordered union must be deterministic across runs") +} + +// --- Q11: score variables under directives ------------------------------------ + +// TestScoreVarDirectiveMatrix exercises a ranker-bound score variable under +// @groupby, @cascade, @normalize, and empty-block aggregation, each checked +// against an oracle derived from a plain val(f) projection of the same variable. +func TestScoreVarDirectiveMatrix(t *testing.T) { + const ( + catPred = "hyb_cat" + tagPred = "hyb_tag" + ) + setSchema(fmt.Sprintf("%s: string .\n%s: string .", catPred, tagPred)) + t.Cleanup(func() { + dropPredicate(catPred) + dropPredicate(tagPred) + }) + // Categories cover all five "fox" matches; tags cover all EXCEPT the top-scored + // doc 503 so @cascade prunes exactly the rank-1 result. Tag values encode the + // decimal uid so @normalize rows can be re-keyed to their uid. + require.NoError(t, addTriplesToCluster(fmt.Sprintf(` + <501> <%[1]s> "groupA" . + <502> <%[1]s> "groupA" . + <503> <%[1]s> "groupB" . + <506> <%[1]s> "groupB" . + <507> <%[1]s> "groupA" . + <501> <%[2]s> "501" . + <502> <%[2]s> "502" . + <506> <%[2]s> "506" . + <507> <%[2]s> "507" . + `, catPred, tagPred))) + + // Oracle: per-uid bm25 scores for the "fox" channel. + oracleJs := processQueryNoErr(t, ` + { + f as var(func: bm25(description_bm25, "fox")) + all(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`) + oracleRows := blockValRows(t, oracleJs, "all", "f") + require.Len(t, oracleRows, 5) + oracle := make(map[string]float64, 5) + for _, r := range oracleRows { + oracle[r.UID] = r.Score + } + require.Equal(t, uidHex(t, 503), oracleRows[0].UID, "503 must be the top bm25 'fox' doc") + + // KNOWN-FAILING (Q11): sum(val(f)) under @groupby evaluates the never-executed val child's empty valueMatrix (query/groupby.go:142 index out of range) — the whole query errors instead of aggregating from the var's uid->score map. + t.Run("GroupbySumOverScoreVar", func(t *testing.T) { + // The parser rejects aggregating an externally-defined value variable inside + // @groupby ("Only aggregator/count functions allowed inside @groupby"), so the + // combination cannot execute — the contract this pins is that it fails CLEANLY + // as a parse error rather than panicking in groupby aggregation + // (query/groupby.go valueMatrix indexing, guarded against regardless). If + // groupby ever learns to aggregate score vars, replace this with the + // per-group sum oracle assertions. + _, err := processQuery(context.Background(), t, fmt.Sprintf(` + { + f as var(func: bm25(description_bm25, "fox")) + grouped(func: uid(f)) @groupby(%s) { + total: sum(val(f)) + } + }`, catPred)) + require.Error(t, err, "score-var aggregation under @groupby is unsupported and must error cleanly") + require.Contains(t, err.Error(), "groupby", "must be the parse-time rejection, not a crash") + }) + + t.Run("CascadeAfterOrdering", func(t *testing.T) { + // @cascade prunes the top-scored doc (503 has no tag); first:2 must then + // return the two best TAGGED docs, still in score order with correct scores + // (no uid-order degradation, no double pagination). + js := processQueryNoErr(t, fmt.Sprintf(` + { + f as var(func: bm25(description_bm25, "fox")) + me(func: uid(f), orderdesc: val(f), first: 2) @cascade { + uid + val(f) + %s + } + }`, tagPred)) + rows := blockValRows(t, js, "me", "f") + require.Len(t, rows, 2, "cascade + first:2 must still fill the page from deeper ranks") + require.NotContains(t, + []string{rows[0].UID, rows[1].UID}, uidHex(t, 503), "untagged 503 must be pruned") + require.GreaterOrEqual(t, rows[0].Score, rows[1].Score, "score order must survive cascade") + for _, r := range rows { + require.Equal(t, oracle[r.UID], r.Score, "cascade must not rebind scores") + } + // The two survivors must be the top-2 by score among the tagged docs + // (multiset comparison so equal-score ties can't flake the assertion). + tagged := []float64{ + oracle[uidHex(t, 501)], oracle[uidHex(t, 502)], + oracle[uidHex(t, 506)], oracle[uidHex(t, 507)], + } + sort.Sort(sort.Reverse(sort.Float64Slice(tagged))) + gotScores := []float64{rows[0].Score, rows[1].Score} + sort.Sort(sort.Reverse(sort.Float64Slice(gotScores))) + require.Equal(t, tagged[:2], gotScores, "cascade page must be the best two tagged docs") + }) + + t.Run("NormalizePairsScoreWithUid", func(t *testing.T) { + // @normalize flattens each entity into one row; the (tag, score) pair in a + // row must belong to the SAME uid (tag values encode the uid). + js := processQueryNoErr(t, fmt.Sprintf(` + { + f as var(func: bm25(description_bm25, "fox")) + me(func: uid(f), orderdesc: val(f)) @normalize { + tag: %s + score: val(f) + } + }`, tagPred)) + var resp struct { + Data struct { + Me []map[string]interface{} `json:"me"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + taggedRows := 0 + for _, row := range resp.Data.Me { + tag, ok := row["tag"].(string) + if !ok { + continue // untagged doc (503) flattens without the tag key + } + taggedRows++ + dec, err := strconv.Atoi(tag) + require.NoError(t, err) + score, ok := row["score"].(float64) + require.True(t, ok, "normalized row for %s must carry its score", tag) + require.Equal(t, oracle[uidHex(t, dec)], score, + "normalized row must pair tag and score of the same uid (%s)", tag) + } + require.Equal(t, 4, taggedRows, "all four tagged docs must produce a flattened row") + }) + + t.Run("EmptyBlockAggregatesFullVarMap", func(t *testing.T) { + // Documented population semantics: an empty block's min/max/sum(val(f)) + // aggregates over the ENTIRE variable map (every uid the ranker bound), + // regardless of what a sibling consuming block filters down to. + js := processQueryNoErr(t, ` + { + f as var(func: bm25(description_bm25, "fox")) + filtered(func: uid(f)) @filter(uid(503)) { uid } + agg() { + mn: min(val(f)) + mx: max(val(f)) + sm: sum(val(f)) + } + }`) + var resp struct { + Data struct { + Agg []map[string]interface{} `json:"agg"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(js), &resp)) + got := make(map[string]float64) + for _, m := range resp.Data.Agg { + for k, v := range m { + if f, ok := v.(float64); ok { + got[k] = f + } + } + } + var mn, mx, sm float64 + mn = math.Inf(1) + for _, s := range oracle { + mn = math.Min(mn, s) + mx = math.Max(mx, s) + sm += s + } + require.InDelta(t, mn, got["mn"], 1e-6, "min must cover the full 5-doc var map") + require.InDelta(t, mx, got["mx"], 1e-6, "max must cover the full 5-doc var map") + require.InDelta(t, sm, got["sm"], 1e-6, + "sum must cover the full var map even though a sibling block filters to one uid") + }) +} + +// --- Q28: degenerate similar_to inputs ---------------------------------------- + +// KNOWN-FAILING (Q28): similar_to with numNeighbors <= 0 reaches the HNSW search with maxResults <= 0 and panics the alpha (tok/hnsw/search_layer.go addPathNode indexes an empty neighbors slice) — no positivity guard, no panic-recovery interceptor. +// Correct behavior asserted here: a clean query-level error (or a defined empty +// result) and a cluster that still serves the next query. NOTE: on today's code +// this test can crash the shared test alpha — run it isolated/last when executing +// against a live cluster. +func TestSimilarToDegenerateK(t *testing.T) { + for _, k := range []string{"0", "-1"} { + t.Run("k="+k, func(t *testing.T) { + query := fmt.Sprintf(` + { + me(func: similar_to(description_vec, %s, "[3.0, 0.0, 0.0, 0.0]")) { uid } + }`, k) + // Either a clean error or a defined (possibly empty) result is acceptable; + // what is NOT acceptable is taking the alpha down. + _, _ = processQuery(context.Background(), t, query) + + // The cluster must still answer a trivial follow-up query. + js := processQueryNoErr(t, `{ me(func: uid(503)) { uid } }`) + require.Contains(t, js, uidHex(t, 503), + "cluster must survive a degenerate similar_to k=%s", k) + }) + } +} + +// TestHybridTopKNonPositiveRejected pins the dql/hybrid.go guard (commit +// dc12fbbc1): hybrid() must reject a non-positive topk up front instead of +// forwarding it to similar_to where it would panic the HNSW search. +func TestHybridTopKNonPositiveRejected(t *testing.T) { + q := ` + { + f as var(func: hybrid(description_bm25, "fox", description_vec, "[3.0, 0.0, 0.0, 0.0]", topk: 0)) + me(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err := processQuery(context.Background(), t, q) + require.Error(t, err) + require.Contains(t, err.Error(), "topk must be a positive integer") + + // A negative topk must also fail (whether at the lexer or the hybrid guard). + qneg := ` + { + f as var(func: hybrid(description_bm25, "fox", description_vec, "[3.0, 0.0, 0.0, 0.0]", topk: -1)) + me(func: uid(f), orderdesc: val(f)) { uid } + }` + _, err = processQuery(context.Background(), t, qneg) + require.Error(t, err) +} + +// --- Q30: fuse-of-fuse and shared channels ------------------------------------ + +// TestFuseOfFuseAndSharedChannels pins scheduler and purity behavior for +// composed fusions in ONE query: a channel consumed by several fuse blocks, a +// fuse block consuming another fuse's output (a -> f1 -> f3 dependency chain), +// and a three-channel fuse. Every fused score is checked against an independent +// RRF oracle computed from the projected raw channel scores — equality proves +// both that the chain resolved (no "Query couldn't be executed") and that no +// fuse mutated a shared channel's score map. +func TestFuseOfFuseAndSharedChannels(t *testing.T) { + js := processQueryNoErr(t, ` + { + a as var(func: bm25(description_bm25, "fox")) + b as var(func: bm25(description_bm25, "dog")) + c as var(func: similar_to(description_vec, 7, "[3.0, 0.0, 0.0, 0.0]")) + f1 as var(func: fuse(a, b, method: "rrf", k: 60)) + f2 as var(func: fuse(a, c, method: "rrf", k: 60, weights: "0.7,0.3")) + f3 as var(func: fuse(f1, c, method: "rrf", k: 60)) + f4 as var(func: fuse(a, b, c, method: "rrf", k: 60)) + cha(func: uid(a)) { uid val(a) } + chb(func: uid(b)) { uid val(b) } + chc(func: uid(c)) { uid val(c) } + r1(func: uid(f1), orderdesc: val(f1)) { uid val(f1) } + r2(func: uid(f2), orderdesc: val(f2)) { uid val(f2) } + r3(func: uid(f3), orderdesc: val(f3)) { uid val(f3) } + r4(func: uid(f4), orderdesc: val(f4)) { uid val(f4) } + }`) + + chA := blockScoreMap(t, js, "cha", "a") + chB := blockScoreMap(t, js, "chb", "b") + chC := blockScoreMap(t, js, "chc", "c") + require.NotEmpty(t, chA) + require.NotEmpty(t, chB) + require.NotEmpty(t, chC) + + requireFusedEquals := func(block, varName string, want map[string]float64) { + got := blockScoreMap(t, js, block, varName) + require.Len(t, got, len(want), "%s: fused uid set must equal the oracle union", block) + for uid, w := range want { + g, ok := got[uid] + require.True(t, ok, "%s: uid %s missing from fused output", block, uid) + require.InDelta(t, w, g, 1e-9, "%s: fused score mismatch for %s", block, uid) + } + } + + // f1 = rrf(a, b); f2 = weighted rrf(a, c) — 'a' is shared by f1/f2/f4, so a + // match here proves f1's computation did not mutate channel a. + f1Want := rrfExpected([]map[string]float64{chA, chB}, nil, 60) + requireFusedEquals("r1", "f1", f1Want) + requireFusedEquals("r2", "f2", + rrfExpected([]map[string]float64{chA, chC}, []float64{0.7, 0.3}, 60)) + + // f3 = rrf(f1, c): fuse output used as a channel (fuse-of-fuse). Compose the + // oracle from the SERVER-reported f1 scores so this check isolates f3. + f1Server := blockScoreMap(t, js, "r1", "f1") + requireFusedEquals("r3", "f3", rrfExpected([]map[string]float64{f1Server, chC}, nil, 60)) + + // Three-channel fuse. + requireFusedEquals("r4", "f4", rrfExpected([]map[string]float64{chA, chB, chC}, nil, 60)) +} diff --git a/schema/parse.go b/schema/parse.go index a18f7fddadb..f8c1d9532cf 100644 --- a/schema/parse.go +++ b/schema/parse.go @@ -443,6 +443,39 @@ func resolveTokenizers(updates []*pb.SchemaUpdate) error { if !has { return errors.Errorf("Invalid tokenizer %s", t) } + if schema.GetList() && tokenizer.Name() == "bm25" { + // BM25 scores a single document (one value) per UID: document length + // and corpus statistics are not well-defined for a list predicate, and + // the bulk loader's per-(attr, uid) stats dedup would silently collapse + // distinct list values. The live path also rejects this combination in + // worker.ValidateSchema; enforcing it here covers every schema ingest + // path (alter, bulk, live loader). + return errors.Errorf( + "Tokenizer 'bm25' cannot be applied to list predicate: %s", + x.ParseAttr(schema.Predicate)) + } + if schema.Lang && tokenizer.Name() == "bm25" { + // Lang-tagged values are indexed under lang-qualified keys and are only + // searchable through a pred@lang qualifier, which bm25() does not + // support. Allowing the combination silently makes every tagged value + // unsearchable (and its corpus-stats contribution asymmetric). Reject it + // until bm25 defines @lang semantics. + return errors.Errorf( + "Tokenizer bm25 cannot be used with @lang on predicate %s: "+ + "bm25 does not support language-qualified values", + x.ParseAttr(schema.Predicate)) + } + if schema.NoConflict && tokenizer.Name() == "bm25" { + // BM25 maintains corpus statistics (doc count, total terms) via a + // read-modify-write that relies on transaction conflict detection to + // serialize concurrent updates to the same stats bucket. @noconflict + // emits no conflict key, so concurrent updates would silently lose + // counts and drift avgDL/IDF. Reject the combination. + return errors.Errorf( + "Tokenizer bm25 cannot be used with @noconflict on predicate %s: "+ + "BM25 corpus statistics require conflict detection to stay consistent", + x.ParseAttr(schema.Predicate)) + } tokenizerType, ok := types.TypeForName(tokenizer.Type()) x.AssertTrue(ok) // Type is validated during tokenizer loading. if tokenizerType != typ { diff --git a/schema/parse_bm25_list_test.go b/schema/parse_bm25_list_test.go new file mode 100644 index 00000000000..fb00a5326b6 --- /dev/null +++ b/schema/parse_bm25_list_test.go @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package schema + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// The bm25 schema invariants must hold at PARSE level so every ingest path is +// covered (alter, live loader, and crucially the bulk loader, which never calls +// worker.ValidateSchema). The bulk stats fold dedupes per (attr, uid), so a list +// predicate slipping through would silently collapse distinct values' stats. +func TestParseRejectsBM25OnListPredicate(t *testing.T) { + _, err := Parse(`pred: [string] @index(bm25) .`) + require.Error(t, err) + require.Contains(t, err.Error(), "cannot be applied to list predicate") +} + +func TestParseRejectsBM25WithLang(t *testing.T) { + _, err := Parse(`pred: string @lang @index(bm25) .`) + require.Error(t, err) + require.Contains(t, err.Error(), "bm25 does not support language-qualified values") +} diff --git a/schema/parse_test.go b/schema/parse_test.go index adb64311f2d..b95f60eb61e 100644 --- a/schema/parse_test.go +++ b/schema/parse_test.go @@ -148,6 +148,20 @@ func TestSchemaIndex_Error1(t *testing.T) { require.Error(t, ParseBytes([]byte(schemaIndexVal2), 1)) } +// BM25 corpus-stats maintenance relies on transaction conflict detection, which +// @noconflict disables — the combination would drift avgDL/IDF and must be rejected. +func TestSchemaBM25NoConflict_Error(t *testing.T) { + err := ParseBytes([]byte(`description: string @index(bm25) @noconflict .`), 1) + require.Error(t, err) + require.Contains(t, err.Error(), "bm25") + require.Contains(t, err.Error(), "noconflict") +} + +// The guard is specific to bm25: a bm25 index without @noconflict is valid. +func TestSchemaBM25WithoutNoConflict(t *testing.T) { + require.NoError(t, ParseBytes([]byte(`description: string @index(bm25) .`), 1)) +} + var schemaIndexVal3Uid = ` person: uid @index . ` diff --git a/systest/bm25lifecycle/multigroup_test.go b/systest/bm25lifecycle/multigroup_test.go new file mode 100644 index 00000000000..28921f11c58 --- /dev/null +++ b/systest/bm25lifecycle/multigroup_test.go @@ -0,0 +1,651 @@ +//go:build integration2 + +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Wave-2 distributed lifecycle tests for native hybrid search (fuse + hybrid). +// +// Q17 TestHybridMultiGroupParity: 3 alphas / replicas=1 (3 groups) with the +// bm25 text predicate and the HNSW vector predicate forced onto DIFFERENT +// groups (zero /moveTablet), versus a 1-alpha control cluster loaded with the +// identical corpus. The same hybrid()/fuse()/single-channel query battery must +// return the identical uid order and scores (1e-9) on both topologies. Any +// divergence localizes to channel content (single-channel legs), not fusion. +// +// Q19 TestBM25ReplicaConvergenceUnderFailure: 3 alphas / replicas=3 (one +// group). Mixed mutation load with a follower stopped mid-load and restarted; +// every replica must then serve identical bm25 rows (each alpha queried +// directly). A >=512-edge single-txn probe then exercises the confirmed +// intra-proposal parallel stats RMW race, and finally the group leader is +// SIGKILLed and post-failover scores must be unchanged. +package main + +import ( + "encoding/json" + "fmt" + "sort" + "strconv" + "strings" + "testing" + "time" + + "github.com/dgraph-io/dgo/v250/protos/api" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/encoding/protojson" + + "github.com/dgraph-io/dgraph/v25/dgraphapi" + "github.com/dgraph-io/dgraph/v25/dgraphtest" + "github.com/dgraph-io/dgraph/v25/protos/pb" +) + +// --------------------------------------------------------------------------- +// Shared helpers (mg*/rc* prefixes to avoid collisions with sibling files in +// this package). +// --------------------------------------------------------------------------- + +// mgRow is one {uid, val(f)} row of a score-ordered result block named "q". +type mgRow struct { + UID string `json:"uid"` + Score float64 `json:"val(f)"` +} + +func mgQueryRows(dg *dgraphapi.GrpcClient, query string) ([]mgRow, error) { + resp, err := dg.Query(query) + if err != nil { + return nil, err + } + var out struct { + Q []mgRow `json:"q"` + } + if err := json.Unmarshal(resp.GetJson(), &out); err != nil { + return nil, err + } + return out.Q, nil +} + +// mgCanonical returns rows sorted by (score desc, uid asc). Score ties are +// possible by construction (RRF gives byte-equal scores to docs holding equal +// ranks in disjoint channels) and the server's sort is not guaranteed stable +// across processes, so all cross-process comparisons are done on this +// canonical form; the per-uid scores and the score-distinct portion of the +// ordering keep full discriminating power. +func mgCanonical(rows []mgRow) []mgRow { + out := make([]mgRow, len(rows)) + copy(out, rows) + sort.SliceStable(out, func(i, j int) bool { + if out[i].Score != out[j].Score { + return out[i].Score > out[j].Score + } + return out[i].UID < out[j].UID + }) + return out +} + +func mgRowsEqual(a, b []mgRow) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].UID != b[i].UID || a[i].Score != b[i].Score { + return false + } + } + return true +} + +func mgSortedScores(rows []mgRow) []float64 { + out := make([]float64, 0, len(rows)) + for _, r := range rows { + out = append(out, r.Score) + } + sort.Sort(sort.Reverse(sort.Float64Slice(out))) + return out +} + +// mgAlphaMembership fetches an alpha's /state and parses the full membership +// snapshot (groups, tablets, members incl. leader flags). +func mgAlphaMembership(hc *dgraphapi.HTTPClient) (*pb.MembershipState, error) { + raw, err := hc.GetAlphaState() + if err != nil { + return nil, err + } + var state pb.MembershipState + if err := protojson.Unmarshal(raw, &state); err != nil { + return nil, err + } + return &state, nil +} + +// mgTabletGroup returns the group serving pred (tablet keys may be plain or +// namespaced like "0-pred"), or 0 if unassigned. +func mgTabletGroup(state *pb.MembershipState, pred string) uint32 { + for gid, group := range state.Groups { + for key := range group.Tablets { + if key == pred || strings.HasSuffix(key, "-"+pred) { + return gid + } + } + } + return 0 +} + +// mgAlphaIdxFromAddr maps a member address like "alpha2:7080" to the +// LocalCluster alpha index (aliasName format is "alpha%d", dgraphtest/dgraph.go). +func mgAlphaIdxFromAddr(addr string) int { + if !strings.HasPrefix(addr, "alpha") { + return -1 + } + numStr := strings.Split(strings.TrimPrefix(addr, "alpha"), ":")[0] + n, err := strconv.Atoi(numStr) + if err != nil { + return -1 + } + return n +} + +func mgLogTablets(t *testing.T, hc *dgraphapi.HTTPClient, label string) { + t.Helper() + state, err := mgAlphaMembership(hc) + if err != nil { + t.Logf("%s: could not fetch membership state: %v", label, err) + return + } + for gid, group := range state.Groups { + for key := range group.Tablets { + t.Logf("%s: group %d serves tablet %q", label, gid, key) + } + } +} + +// --------------------------------------------------------------------------- +// Q17: multi-group parity for hybrid()/fuse(). +// --------------------------------------------------------------------------- + +const mgSchema = ` + mgp_text: string @index(bm25) . + mgp_vec: float32vector @index(hnsw(metric:"euclidean")) . +` + +// mgCorpusNquads builds 96 docs (3 full uid%32 bucket cycles) with explicit +// uids so both clusters hold byte-identical data. Text has varying "fox"/"dog" +// term frequencies and doc lengths; vectors are distinct 4-dim points. +func mgCorpusNquads() string { + var b strings.Builder + for i := 0; i < 96; i++ { + uid := 0x1000 + i + text := strings.TrimSpace( + strings.Repeat("fox ", i%5+1) + strings.Repeat("dog ", i%3) + fmt.Sprintf("word%d", i%11)) + fmt.Fprintf(&b, "<%#x> %q .\n", uid, text) + fmt.Fprintf(&b, "<%#x> \"[%d.0, %d.0, %d.0, 1.0]\" .\n", uid, i%13, (i*7)%17, i%5) + } + return b.String() +} + +// mgParityBattery is the identical query set run against both topologies. The +// single-channel legs localize any divergence to a channel; the fuse/hybrid +// legs assert coordinator-side fusion parity on top. +var mgParityBattery = []struct { + name string + // exact: demand bit-identical cross-topology parity. Only the bm25 leg is + // exact — it is a deterministic algorithm over identical data. Legs touching + // similar_to are APPROXIMATE: the two topologies build independent HNSW + // graphs (different insertion interleaving), which legitimately differ on + // boundary-rank candidates. For those legs we assert the invariants the + // product does promise: row counts, score-descending order, >= n-1 uid-set + // overlap (one boundary swap allowed), and — the misbinding detector — + // bit-equal scores for every uid both topologies return. + exact bool + // scoreCheck: per-uid scores for uids BOTH topologies return must be equal. + // True for single-channel legs (bm25 is deterministic; a similarity score is + // a pure function of stored vector + query vector, so equality holds even + // when the approximate graphs disagree on boundary candidates). False for + // fused legs: fused scores inherit rank noise from the approximate vector + // channel (a boundary uid in one topology's top-k but not the other's shifts + // every RRF contribution), and per-uid fusion correctness is already pinned + // by the single-cluster oracle test in query/query_hybrid_test.go. + scoreCheck bool + query string +}{ + {"bm25_channel_only", true, true, `{ + f as var(func: bm25(mgp_text, "fox dog")) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`}, + {"vector_channel_only", false, true, `{ + f as var(func: similar_to(mgp_vec, 12, "[1.0, 2.0, 3.0, 4.0]")) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`}, + {"fuse_rrf", false, false, `{ + txt as var(func: bm25(mgp_text, "fox dog")) + vec as var(func: similar_to(mgp_vec, 12, "[1.0, 2.0, 3.0, 4.0]")) + f as var(func: fuse(txt, vec, method: "rrf", k: 60)) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`}, + {"fuse_linear", false, false, `{ + txt as var(func: bm25(mgp_text, "fox dog")) + vec as var(func: similar_to(mgp_vec, 12, "[1.0, 2.0, 3.0, 4.0]")) + f as var(func: fuse(txt, vec, method: "linear", weights: "0.3,0.7", normalize: "max")) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`}, + {"hybrid_sugar", false, false, `{ + f as var(func: hybrid(mgp_text, "fox dog", mgp_vec, "[1.0, 2.0, 3.0, 4.0]", topk: 12, method: "rrf", k: 60)) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } + }`}, +} + +func TestHybridMultiGroupParity(t *testing.T) { + // 3 alphas x replicas=1 => three single-member groups. + conf3 := dgraphtest.NewClusterConfig().WithNumAlphas(3).WithNumZeros(1).WithReplicas(1) + c3, err := dgraphtest.NewLocalCluster(conf3) + require.NoError(t, err) + defer func() { c3.Cleanup(t.Failed()) }() + require.NoError(t, c3.Start()) + + // 1-alpha control cluster: same data, trivially colocated tablets. + conf1 := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1).WithReplicas(1) + c1, err := dgraphtest.NewLocalCluster(conf1) + require.NoError(t, err) + defer func() { c1.Cleanup(t.Failed()) }() + require.NoError(t, c1.Start()) + + dg3, cleanup3, err := c3.Client() + require.NoError(t, err) + defer cleanup3() + dg1, cleanup1, err := c1.Client() + require.NoError(t, err) + defer cleanup1() + + require.NoError(t, dg3.DropAll()) + require.NoError(t, dg1.DropAll()) + require.NoError(t, dg3.SetupSchema(mgSchema)) + require.NoError(t, dg1.SetupSchema(mgSchema)) + time.Sleep(2 * time.Second) // post-alter settling, as existing systest does + + hc3, err := c3.HTTPClient() + require.NoError(t, err) + + // Force mgp_text and mgp_vec onto DIFFERENT groups before any data lands. + // The move attempt lives inside the poll: retried until zero's state shows + // the split (MoveTablet is a no-op error once the tablet already moved). + var textGid, vecGid uint32 + require.Eventually(t, func() bool { + state, err := mgAlphaMembership(hc3) + if err != nil { + return false + } + textGid = mgTabletGroup(state, "mgp_text") + vecGid = mgTabletGroup(state, "mgp_vec") + if textGid == 0 || vecGid == 0 { + return false + } + if textGid != vecGid { + return true + } + var target uint32 + for gid := range state.Groups { + if gid != 0 && gid != textGid { + target = gid + break + } + } + if target == 0 { + return false + } + if err := hc3.MoveTablet("mgp_vec", target); err != nil { + t.Logf("moveTablet(mgp_vec -> group %d): %v (will re-check)", target, err) + } + return false + }, 60*time.Second, 2*time.Second, + "mgp_text and mgp_vec must end up on different groups") + t.Logf("placement: mgp_text on group %d, mgp_vec on group %d", textGid, vecGid) + + // Identical corpus into both clusters. 192 edges per txn (< 512) keeps the + // proposal application serial, so this load is not subject to the Q19 race. + corpus := mgCorpusNquads() + _, err = dg3.Mutate(&api.Mutation{SetNquads: []byte(corpus), CommitNow: true}) + require.NoError(t, err) + _, err = dg1.Mutate(&api.Mutation{SetNquads: []byte(corpus), CommitNow: true}) + require.NoError(t, err) + + // Observe where the HNSW auxiliary __vector_* tablets landed (Q17's + // statically undecidable hazard). Logged for diagnosis; the parity battery + // below is the load-bearing assertion. + mgLogTablets(t, hc3, "3-group cluster post-load") + + // Wait until the (possibly just-moved) vector tablet is queryable on both + // clusters: membership propagation to the alphas can lag the move, and + // ErrNonExistentTablet silently yields empty results in the interim. + vecProbe := mgParityBattery[1].query + require.Eventually(t, func() bool { + r3, err3 := mgQueryRows(dg3, vecProbe) + r1, err1 := mgQueryRows(dg1, vecProbe) + return err3 == nil && err1 == nil && len(r3) == 12 && len(r1) == 12 + }, 60*time.Second, 2*time.Second, + "similar_to must return topk results on both clusters before the parity battery") + + for _, tc := range mgParityBattery { + t.Run(tc.name, func(t *testing.T) { + rows3, err := mgQueryRows(dg3, tc.query) + require.NoError(t, err) + rows1, err := mgQueryRows(dg1, tc.query) + require.NoError(t, err) + require.NotEmpty(t, rows1, "control cluster returned no rows") + require.NotEmpty(t, rows3, "3-group cluster returned no rows") + + // Server-side ordering must be score-descending on both. + for i := 1; i < len(rows3); i++ { + require.GreaterOrEqual(t, rows3[i-1].Score, rows3[i].Score, + "3-group rows not score-descending at index %d", i) + } + for i := 1; i < len(rows1); i++ { + require.GreaterOrEqual(t, rows1[i-1].Score, rows1[i].Score, + "control rows not score-descending at index %d", i) + } + + want := mgCanonical(rows1) + got := mgCanonical(rows3) + require.Equal(t, len(want), len(got), + "row count differs between topologies") + + if tc.exact { + // Deterministic leg: identical uid order and scores (1e-9), on the + // canonical form (score desc, uid asc) so exact-tie permutations + // cannot flake. + for i := range want { + require.Equalf(t, want[i].UID, got[i].UID, + "uid order diverges at rank %d (control %s / 3-group %s)", + i, want[i].UID, got[i].UID) + require.InDeltaf(t, want[i].Score, got[i].Score, 1e-9, + "score diverges for uid %s at rank %d", want[i].UID, i) + } + return + } + + // Approximate leg (touches HNSW): allow one boundary-rank swap between + // the independently built graphs, but scores for uids BOTH return must + // be bit-close — a mismatch there is score misbinding, not + // approximation. + wantScores := make(map[string]float64, len(want)) + for _, r := range want { + wantScores[r.UID] = r.Score + } + shared := 0 + for _, r := range got { + if ws, ok := wantScores[r.UID]; ok { + shared++ + if tc.scoreCheck { + require.InDeltaf(t, ws, r.Score, 1e-9, + "score for uid %s differs between topologies (misbinding)", r.UID) + } + } + } + require.GreaterOrEqualf(t, shared, len(want)-1, + "uid overlap %d/%d below n-1: more than a boundary-rank divergence", + shared, len(want)) + }) + } +} + +// --------------------------------------------------------------------------- +// Q19: single-group replica convergence under follower outage, heavy-txn +// stats race, and leader failover. +// --------------------------------------------------------------------------- + +func rcDocNquads(start, count int) string { + var b strings.Builder + for i := start; i < start+count; i++ { + uid := 0x2000 + i + text := strings.TrimSpace( + strings.Repeat("zephyr ", i%4+1) + strings.Repeat("quill ", i%3) + fmt.Sprintf("tag%d", i%9)) + fmt.Fprintf(&b, "<%#x> %q .\n", uid, text) + } + return b.String() +} + +func rcMutateBatches(dg *dgraphapi.GrpcClient, startDoc, batches, perBatch int) error { + for b := 0; b < batches; b++ { + nq := rcDocNquads(startDoc+b*perBatch, perBatch) + if _, err := dg.Mutate(&api.Mutation{SetNquads: []byte(nq), CommitNow: true}); err != nil { + return err + } + } + return nil +} + +func rcRankQuery(pred, terms string) string { + return fmt.Sprintf(`{ + f as var(func: bm25(%s, %q)) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } +}`, pred, terms) +} + +const rcScoreQuery = `{ + f as var(func: bm25(rc_text, "zephyr quill")) + q(func: uid(f), orderdesc: val(f)) { uid val(f) } +}` + +func TestBM25ReplicaConvergenceUnderFailure(t *testing.T) { + // 3 alphas x replicas=3 => one group, three replicas. + conf := dgraphtest.NewClusterConfig().WithNumAlphas(3).WithNumZeros(1).WithReplicas(3) + c, err := dgraphtest.NewLocalCluster(conf) + require.NoError(t, err) + defer func() { c.Cleanup(t.Failed()) }() + require.NoError(t, c.Start()) + + dg, cleanupAll, err := c.Client() + require.NoError(t, err) + defer cleanupAll() + + require.NoError(t, dg.DropAll()) + // Explicit uids go up to 0xA000+1024 (~42k) in the heavy-txn subtest; zero's + // default lease is far below that, so extend it up front. + require.NoError(t, c.AssignUids(dg.Dgraph, 50000)) + require.NoError(t, dg.SetupSchema(`rc_text: string @index(bm25) .`)) + time.Sleep(2 * time.Second) // post-alter settling + + // Phase A: 8 x 20-doc txns (docs 0..159) with all replicas up. Every txn + // stays far below the 512-edge parallel-apply threshold, so replicas apply + // stats strictly serially and MUST agree. + require.NoError(t, rcMutateBatches(dg, 0, 8, 20)) + + // Identify the group leader and the two followers. + hc, err := c.HTTPClient() + require.NoError(t, err) + leaderIdx := -1 + var followerIdxs []int + require.Eventually(t, func() bool { + state, err := mgAlphaMembership(hc) + if err != nil { + return false + } + leaderIdx = -1 + followerIdxs = followerIdxs[:0] + for _, group := range state.Groups { + for _, m := range group.Members { + idx := mgAlphaIdxFromAddr(m.Addr) + if idx < 0 { + return false + } + if m.Leader { + if leaderIdx >= 0 { + return false // stale double-leader view; re-poll + } + leaderIdx = idx + } else { + followerIdxs = append(followerIdxs, idx) + } + } + } + return leaderIdx >= 0 && len(followerIdxs) == 2 + }, 60*time.Second, 2*time.Second, "must observe one leader and two followers") + sort.Ints(followerIdxs) + stopIdx, pinIdx := followerIdxs[0], followerIdxs[1] + t.Logf("leader alpha %d; stopping follower alpha %d; pinning writes to follower alpha %d", + leaderIdx, stopIdx, pinIdx) + + // dgPin targets an alpha that is never stopped or killed, so the mixed load + // keeps flowing while a follower is down and after the leader dies. + dgPin, cleanupPin, err := c.AlphaClient(pinIdx) + require.NoError(t, err) + defer cleanupPin() + + // Phase B: stop a follower mid-load; insert docs 160..319 and rewrite the + // text (paired DEL+SET stats path) of docs 0..19 while it is down. + require.NoError(t, c.StopAlpha(stopIdx)) + require.NoError(t, rcMutateBatches(dgPin, 160, 8, 20)) + var upd strings.Builder + for i := 0; i < 20; i++ { + fmt.Fprintf(&upd, "<%#x> %q .\n", 0x2000+i, + fmt.Sprintf("quill quill zephyr renumbered%d", i)) + } + _, err = dgPin.Mutate(&api.Mutation{SetNquads: []byte(upd.String()), CommitNow: true}) + require.NoError(t, err) + + // Phase C: restart the follower, then keep mutating: insert docs 320..479 + // and delete docs 20..29 entirely (S P * on an un-lang-tagged predicate). + require.NoError(t, c.StartAlpha(stopIdx)) + require.NoError(t, rcMutateBatches(dgPin, 320, 8, 20)) + var del strings.Builder + for i := 20; i < 30; i++ { + fmt.Fprintf(&del, "<%#x> * .\n", 0x2000+i) + } + _, err = dgPin.Mutate(&api.Mutation{DelNquads: []byte(del.String()), CommitNow: true}) + require.NoError(t, err) + + // Reference ranking from the pinned alpha after the load quiesces. + // 470 live docs all match "zephyr". + var ref []mgRow + require.Eventually(t, func() bool { + rows, err := mgQueryRows(dgPin, rcScoreQuery) + if err != nil || len(rows) != 470 { + return false + } + ref = mgCanonical(rows) + return true + }, 60*time.Second, 2*time.Second, "pinned alpha must serve the full corpus") + + // Q19 assertion 1: every replica — including the stopped-and-restarted + // follower, which must catch up over Raft — serves the identical ranking + // with identical (bitwise) scores. Expected to PASS: all txns above were + // below the parallel-apply threshold. + t.Run("AllReplicasIdenticalAfterCatchup", func(t *testing.T) { + for idx := 0; idx < 3; idx++ { + require.Eventuallyf(t, func() bool { + dgi, cl, err := c.AlphaClient(idx) + if err != nil { + return false + } + defer cl() + rows, err := mgQueryRows(dgi, rcScoreQuery) + if err != nil { + return false + } + return mgRowsEqual(ref, mgCanonical(rows)) + }, 120*time.Second, 2*time.Second, + "alpha %d must converge to the reference bm25 ranking", idx) + } + }) + + // Q19 assertion 2 (confirmed-vulnerable probe): one >=512-edge proposal + // splits into concurrent apply goroutines whose unlocked stats + // read-modify-writes on shared uid%32 buckets can lose contributions + // (worker/draft.go:561-601 + posting/bm25.go:236-263). Oracle: the same + // 1024 texts loaded as 32-doc txns must yield the identical sorted score + // sequence; replicas must also agree among themselves. EXPECTED TO FAIL + // (possibly intermittently — the race is scheduling-dependent). + t.Run("HeavySingleTxnStatsOracle", func(t *testing.T) { + require.NoError(t, dgPin.SetupSchema(` + rc_heavy: string @index(bm25) . + rc_heavy_ctl: string @index(bm25) .`)) + time.Sleep(2 * time.Second) // post-alter settling + + var heavy, ctl strings.Builder + for i := 0; i < 1024; i++ { + text := strings.TrimSpace( + strings.Repeat("gryphon ", i%6+1) + fmt.Sprintf("mark%d", i%13)) + fmt.Fprintf(&heavy, "<%#x> %q .\n", 0x8000+i, text) + fmt.Fprintf(&ctl, "<%#x> %q .\n", 0xA000+i, text) + } + // Heavy: 1024 edges in ONE txn => parallel apply on every replica. + _, err := dgPin.Mutate(&api.Mutation{SetNquads: []byte(heavy.String()), CommitNow: true}) + require.NoError(t, err) + // Control: identical texts (same uid%32 bucket distribution) in 32-doc + // txns => serial apply, exact stats. + ctlLines := strings.Split(strings.TrimSpace(ctl.String()), "\n") + for i := 0; i < len(ctlLines); i += 32 { + end := i + 32 + if end > len(ctlLines) { + end = len(ctlLines) + } + _, err := dgPin.Mutate(&api.Mutation{ + SetNquads: []byte(strings.Join(ctlLines[i:end], "\n")), + CommitNow: true, + }) + require.NoError(t, err) + } + + heavyQ := rcRankQuery("rc_heavy", "gryphon") + ctlQ := rcRankQuery("rc_heavy_ctl", "gryphon") + var refHeavy []mgRow + for idx := 0; idx < 3; idx++ { + var heavyRows, ctlRows []mgRow + require.Eventuallyf(t, func() bool { + dgi, cl, err := c.AlphaClient(idx) + if err != nil { + return false + } + defer cl() + h, err := mgQueryRows(dgi, heavyQ) + if err != nil || len(h) != 1024 { + return false + } + cr, err := mgQueryRows(dgi, ctlQ) + if err != nil || len(cr) != 1024 { + return false + } + heavyRows, ctlRows = h, cr + return true + }, 90*time.Second, 2*time.Second, + "alpha %d must serve both heavy and control corpora", idx) + + // Within-alpha oracle: identical corpora => identical score sets. + // A lost stats RMW shifts N/avgDL and therefore every score. + require.Equalf(t, mgSortedScores(ctlRows), mgSortedScores(heavyRows), + "alpha %d: 1024-doc single-txn corpus scores differ from the serially-loaded "+ + "identical corpus — intra-proposal parallel stats RMW lost an update (Q19/Q15)", idx) + + // Cross-replica identity for the heavy corpus. + canon := mgCanonical(heavyRows) + if idx == 0 { + refHeavy = canon + } else { + require.Truef(t, mgRowsEqual(refHeavy, canon), + "alpha %d heavy-corpus ranking diverges from alpha 0 — replicas resolved "+ + "the stats race differently (Q19)", idx) + } + } + }) + + // Q19 assertion 3: SIGKILL the group leader; after failover the two + // surviving replicas must serve the unchanged reference ranking. Expected + // to PASS. + t.Run("LeaderFailoverScoresUnchanged", func(t *testing.T) { + require.NoError(t, c.KillAlpha(leaderIdx)) + for _, idx := range []int{stopIdx, pinIdx} { + require.Eventuallyf(t, func() bool { + dgi, cl, err := c.AlphaClient(idx) + if err != nil { + return false + } + defer cl() + rows, err := mgQueryRows(dgi, rcScoreQuery) + if err != nil { + return false + } + return mgRowsEqual(ref, mgCanonical(rows)) + }, 120*time.Second, 2*time.Second, + "alpha %d must serve the unchanged ranking after leader failover", idx) + } + }) +} diff --git a/systest/bm25lifecycle/parity_test.go b/systest/bm25lifecycle/parity_test.go new file mode 100644 index 00000000000..70ef582153e --- /dev/null +++ b/systest/bm25lifecycle/parity_test.go @@ -0,0 +1,624 @@ +//go:build integration2 + +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package main + +import ( + "encoding/json" + "fmt" + "math" + "math/rand" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "testing" + "time" + + "github.com/dgraph-io/dgo/v250/protos/api" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/dgraphapi" + "github.com/dgraph-io/dgraph/v25/dgraphtest" + "github.com/dgraph-io/dgraph/v25/tok" +) + +// --------------------------------------------------------------------------- +// Shared helpers +// --------------------------------------------------------------------------- + +type scoredUID struct { + uid uint64 + score float64 +} + +func parseScored(jsonBytes []byte) (map[uint64]float64, error) { + var out struct { + Q []struct { + UID string `json:"uid"` + Score float64 `json:"score"` + } `json:"q"` + } + if err := json.Unmarshal(jsonBytes, &out); err != nil { + return nil, err + } + res := make(map[uint64]float64, len(out.Q)) + for _, row := range out.Q { + uid, err := strconv.ParseUint(strings.TrimPrefix(row.UID, "0x"), 16, 64) + if err != nil { + return nil, err + } + res[uid] = row.Score + } + return res, nil +} + +// --------------------------------------------------------------------------- +// Q23: bulk vs live parity for BM25 (bulk mapper counts stats per nquad while +// reduce dedupes postings, so duplicated triples inflate docCount/totalTerms +// on the bulk path only). +// --------------------------------------------------------------------------- + +const ( + parityPred = "parity_text" + parityNDocs = 1000 + parityK = 1.2 + parityB = 0.75 + parityProbe = "quixotic" + parityProbe1 = uint64(101) + parityProbe2 = uint64(507) +) + +type parityDoc struct { + uid uint64 + text string + tf map[string]uint32 + dl uint32 +} + +type parityCorpus struct { + docs map[uint64]*parityDoc + order []uint64 // ascending uids + n float64 // docCount (each doc counted once) + totalTerms float64 + df map[string]uint32 +} + +// buildParityCorpus deterministically generates the corpus and computes the +// exact ground-truth stats using the SAME tokenizer the server uses at index +// time (tok.BM25Tokenizer with lang ""), so the oracle is exact by construction. +func buildParityCorpus(t *testing.T) *parityCorpus { + vocab := []string{ + "amber", "basalt", "cobalt", "dune", "ember", "fjord", "garnet", "harbor", + "iris", "juniper", "krypton", "lagoon", "marble", "nectar", "onyx", "pumice", + "quartz", "russet", "sable", "topaz", "umber", "velvet", "willow", "zephyr", + } + r := rand.New(rand.NewSource(42)) + c := &parityCorpus{ + docs: make(map[uint64]*parityDoc, parityNDocs), + df: make(map[string]uint32), + } + bt := tok.BM25Tokenizer{} + for i := 1; i <= parityNDocs; i++ { + uid := uint64(i) + n := 3 + i%10 + words := make([]string, 0, n+1) + for j := 0; j < n; j++ { + words = append(words, vocab[r.Intn(len(vocab))]) + } + if uid == parityProbe1 || uid == parityProbe2 { + words = append(words, parityProbe) + } + text := strings.Join(words, " ") + tf, dl, err := bt.TokensWithFrequency(text, "") + require.NoError(t, err) + require.NotZero(t, dl, "corpus doc %d tokenized to zero terms", i) + c.docs[uid] = &parityDoc{uid: uid, text: text, tf: tf, dl: dl} + c.order = append(c.order, uid) + c.n++ + c.totalTerms += float64(dl) + for term := range tf { + c.df[term]++ + } + } + return c +} + +func parityQueryTokens(t *testing.T, queryText string) []string { + bt := tok.BM25Tokenizer{} + all, err := bt.Tokens(queryText) + require.NoError(t, err) + seen := make(map[string]struct{}, len(all)) + var uniq []string + for _, tkn := range all { + if _, ok := seen[tkn]; !ok { + seen[tkn] = struct{}{} + uniq = append(uniq, tkn) + } + } + return uniq +} + +// bm25TermScore mirrors worker/bm25wand.go bm25Score exactly. +func bm25TermScore(idf, tf, dl, avgDL, k, b float64) float64 { + if avgDL <= 0 { + avgDL = 1 + } + if dl <= 0 { + dl = 1 + } + return idf * (k + 1) * tf / (k*(1-b+b*dl/avgDL) + tf) +} + +// expectedBM25 computes the closed-form expected score for every matching doc. +func (c *parityCorpus) expectedBM25(t *testing.T, queryText string) map[uint64]float64 { + tokens := parityQueryTokens(t, queryText) + avgDL := c.totalTerms / c.n + res := make(map[uint64]float64) + for _, term := range tokens { + df := float64(c.df[term]) + if df == 0 { + continue + } + n := c.n + if n < df { + n = df + } + idf := math.Log1p((n - df + 0.5) / (df + 0.5)) + for uid, doc := range c.docs { + tf := float64(doc.tf[term]) + if tf == 0 { + continue + } + res[uid] += bm25TermScore(idf, tf, float64(doc.dl), avgDL, parityK, parityB) + } + } + return res +} + +func (c *parityCorpus) nquadLines(dupUids map[uint64]bool) []string { + var lines []string + for _, uid := range c.order { + line := fmt.Sprintf("<0x%x> <%s> %q .", uid, parityPred, c.docs[uid].text) + lines = append(lines, line) + if dupUids[uid] { + lines = append(lines, line) + } + } + return lines +} + +func fetchBM25Scores(dg *dgraphapi.GrpcClient, queryText string) (map[uint64]float64, error) { + q := fmt.Sprintf(`{ + s as var(func: bm25(%s, %q)) + q(func: uid(s)) { uid score: val(s) } + }`, parityPred, queryText) + resp, err := dg.Query(q) + if err != nil { + return nil, err + } + return parseScored(resp.GetJson()) +} + +func requireScoreMapsEqual(t *testing.T, label string, want, got map[uint64]float64, delta float64) { + t.Helper() + require.Len(t, got, len(want), + "%s: result count mismatch (want %d docs, got %d)", label, len(want), len(got)) + for uid, w := range want { + g, ok := got[uid] + require.True(t, ok, "%s: uid 0x%x missing from results", label, uid) + require.InDelta(t, w, g, delta, "%s: score mismatch for uid 0x%x", label, uid) + } +} + +// deriveBM25Stats inverts two observed scores of one term (tf, dl known per +// doc, df known) back into the (docCount, avgDL) the server must have used. +// +// s_i = idf * A_i / (E_i + F_i * x) with x = 1/avgDL, +// E_i = k(1-b) + tf_i, F_i = k*b*dl_i, A_i = (k+1)*tf_i, +// idf = ln(1 + (N - df + 0.5)/(df + 0.5)). +func deriveBM25Stats(s1, tf1, dl1, s2, tf2, dl2, df float64) (docCount, avgDL float64) { + k, b := parityK, parityB + e1, f1, a1 := k*(1-b)+tf1, k*b*dl1, (k+1)*tf1 + e2, f2, a2 := k*(1-b)+tf2, k*b*dl2, (k+1)*tf2 + x := (s2*e2/a2 - s1*e1/a1) / (s1*f1/a1 - s2*f2/a2) + avgDL = 1 / x + idf := s1 * (e1 + f1*x) / a1 + docCount = (math.Exp(idf)-1)*(df+0.5) + df - 0.5 + return docCount, avgDL +} + +// TestBM25BulkLiveParity loads the same logical dataset (1000 docs; the RDF +// file duplicates the triples of uids 1..50; the live path additionally +// re-SETs those triples and delete+reinserts uids 901..910) into a bulk-loaded +// cluster and a live-mutated cluster, then asserts: +// 1. the live cluster's scores match the in-test closed-form oracle to 1e-9, +// 2. the bulk cluster's scores match the same oracle (regression pin: the bulk +// mapper used to count stats once per nquad while reduce deduped postings, +// inflating docCount by the duplicate count — fixed by per-doc stats +// postings folded in the reducer), +// 3. both clusters' score-annotated result lists agree to 1e-9, +// 4. (docCount, avgDL) derived closed-form from the probe-term scores equal +// ground truth on both clusters. +func TestBM25BulkLiveParity(t *testing.T) { + corpus := buildParityCorpus(t) + // Q23 regression (fixed): the bulk mapper used to count stats once per nquad + // while reduce deduped postings, so duplicate triples inflated docCount + // (empirically 1050 for truth 1000). The mapper now emits per-document stats + // postings on the bucket keys and the reducer folds them after the same-uid + // dedup, so duplicates cannot inflate stats. Duplicates are therefore part of + // the DEFAULT corpus here, pinning the fix. + dupUids := make(map[uint64]bool) + for i := uint64(1); i <= 50; i++ { + dupUids[i] = true + } + + battery := []string{ + parityProbe, + "amber", + "basalt cobalt", + "ember fjord garnet harbor", + "zephyr quixotic dune", + } + + schemaStr := fmt.Sprintf("%s: string @index(bm25) .", parityPred) + + // --- Cluster A: bulk load an RDF file containing duplicate SET triples. + baseDir := t.TempDir() + rdfFile := filepath.Join(baseDir, "parity.rdf") + require.NoError(t, os.WriteFile(rdfFile, + []byte(strings.Join(corpus.nquadLines(dupUids), "\n")+"\n"), 0o644)) + schemaFile := filepath.Join(baseDir, "parity.schema") + require.NoError(t, os.WriteFile(schemaFile, []byte(schemaStr+"\n"), 0o644)) + + confA := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1). + WithReplicas(1).WithBulkLoadOutDir(t.TempDir()) + bulkC, err := dgraphtest.NewLocalCluster(confA) + require.NoError(t, err) + defer func() { bulkC.Cleanup(t.Failed()) }() + + require.NoError(t, bulkC.StartZero(0)) + require.NoError(t, bulkC.HealthCheck(true)) + require.NoError(t, bulkC.BulkLoad(dgraphtest.BulkOpts{ + DataFiles: []string{rdfFile}, + SchemaFiles: []string{schemaFile}, + })) + require.NoError(t, bulkC.Start()) + + bulkDg, bulkCleanup, err := bulkC.Client() + require.NoError(t, err) + defer bulkCleanup() + + // --- Cluster B: live client mutations for the identical logical dataset. + confB := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1).WithReplicas(1) + liveC, err := dgraphtest.NewLocalCluster(confB) + require.NoError(t, err) + defer func() { liveC.Cleanup(t.Failed()) }() + require.NoError(t, liveC.Start()) + + liveDg, liveCleanup, err := liveC.Client() + require.NoError(t, err) + defer liveCleanup() + + require.NoError(t, liveDg.DropAll()) + require.NoError(t, liveDg.SetupSchema(schemaStr)) + require.NoError(t, liveC.AssignUids(liveDg.Dgraph, 2*parityNDocs)) + + // Initial load in small batches (100 edges/txn) so the known intra-proposal + // parallel stats race (Q15/Q19, >=512 edges per proposal) cannot pollute + // this parity comparison. + baseLines := corpus.nquadLines(nil) + const batchSize = 100 + for start := 0; start < len(baseLines); start += batchSize { + end := start + batchSize + if end > len(baseLines) { + end = len(baseLines) + } + _, err := liveDg.Mutate(&api.Mutation{ + SetNquads: []byte(strings.Join(baseLines[start:end], "\n")), + CommitNow: true, + }) + require.NoError(t, err) + } + + // Duplicate SETs: re-SET the identical triples of the dup'd uids in a + // later txn. The live DEL+SET pairing must net exactly zero stats change. + var dupLines []string + for uid := range dupUids { + dupLines = append(dupLines, fmt.Sprintf("<0x%x> <%s> %q .", uid, parityPred, corpus.docs[uid].text)) + } + _, err = liveDg.Mutate(&api.Mutation{ + SetNquads: []byte(strings.Join(dupLines, "\n")), CommitNow: true, + }) + require.NoError(t, err) + + // Delete + reinsert: wipe and restore uids 901..910 one txn at a time. + for uid := uint64(901); uid <= 910; uid++ { + _, err = liveDg.Mutate(&api.Mutation{ + DelNquads: []byte(fmt.Sprintf("<0x%x> <%s> * .", uid, parityPred)), + CommitNow: true, + }) + require.NoError(t, err) + _, err = liveDg.Mutate(&api.Mutation{ + SetNquads: []byte(fmt.Sprintf("<0x%x> <%s> %q .", uid, parityPred, corpus.docs[uid].text)), + CommitNow: true, + }) + require.NoError(t, err) + } + + // Wait until both clusters answer bm25 queries at all (startup settling). + for _, dg := range []*dgraphapi.GrpcClient{liveDg, bulkDg} { + dg := dg + require.EventuallyWithT(t, func(ct *assert.CollectT) { + got, err := fetchBM25Scores(dg, "amber") + if !assert.NoError(ct, err) { + return + } + assert.NotEmpty(ct, got) + }, 60*time.Second, time.Second) + } + + // 1. Live cluster vs closed-form oracle (expected to pass: pins the + // re-SET / delete+reinsert stats symmetry on the live path). + liveResults := make(map[string]map[uint64]float64, len(battery)) + for _, q := range battery { + want := corpus.expectedBM25(t, q) + got, err := fetchBM25Scores(liveDg, q) + require.NoError(t, err) + requireScoreMapsEqual(t, "live["+q+"]", want, got, 1e-9) + liveResults[q] = got + } + + // 2. Derived (docCount, avgDL) from the probe term on the live cluster. + probeToken := parityQueryTokens(t, parityProbe) + require.Len(t, probeToken, 1) + require.EqualValues(t, 2, corpus.df[probeToken[0]], "probe term must have df=2") + d1, d2 := corpus.docs[parityProbe1], corpus.docs[parityProbe2] + require.NotEqual(t, d1.dl, d2.dl, "probe docs must have distinct doc lengths") + liveProbe := liveResults[parityProbe] + n, avgDL := deriveBM25Stats( + liveProbe[parityProbe1], float64(d1.tf[probeToken[0]]), float64(d1.dl), + liveProbe[parityProbe2], float64(d2.tf[probeToken[0]]), float64(d2.dl), + float64(corpus.df[probeToken[0]])) + require.InDelta(t, corpus.n, n, 0.9, + "live cluster: derived docCount diverges from ground truth") + require.InEpsilon(t, corpus.totalTerms/corpus.n, avgDL, 1e-6, + "live cluster: derived avgDL diverges from ground truth") + + // 3. Bulk cluster vs the same oracle. EXPECTED TO FAIL TODAY (Q23): the + // bulk mapper increments docCount/totalTerms once per nquad while the + // reducer dedupes (key,uid) postings, so the 50 duplicated triples leave + // the bulk cluster with docCount=1050 — every IDF, hence every score, + // deviates from the live cluster and the oracle. + bulkResults := make(map[string]map[uint64]float64, len(battery)) + for _, q := range battery { + got, err := fetchBM25Scores(bulkDg, q) + require.NoError(t, err) + bulkResults[q] = got + } + bulkProbe := bulkResults[parityProbe] + if s1, ok1 := bulkProbe[parityProbe1]; ok1 { + if s2, ok2 := bulkProbe[parityProbe2]; ok2 { + bn, bAvg := deriveBM25Stats( + s1, float64(d1.tf[probeToken[0]]), float64(d1.dl), + s2, float64(d2.tf[probeToken[0]]), float64(d2.dl), + float64(corpus.df[probeToken[0]])) + t.Logf("bulk cluster implied stats: docCount=%.3f avgDL=%.6f (truth: %.0f / %.6f)", + bn, bAvg, corpus.n, corpus.totalTerms/corpus.n) + } + } + for _, q := range battery { + requireScoreMapsEqual(t, "bulk["+q+"]", corpus.expectedBM25(t, q), bulkResults[q], 1e-9) + } + + // 4. Cross-cluster parity: score-annotated result lists equal to 1e-9. + for _, q := range battery { + requireScoreMapsEqual(t, "bulk-vs-live["+q+"]", liveResults[q], bulkResults[q], 1e-9) + } +} + +// --------------------------------------------------------------------------- +// Q22: HNSW lifecycle — upsert flips the winner, delete the winner, reinsert +// the same uid with a different vector, kill+restart the alpha. After every +// step the similar_to top-k (uids AND scores) must match an exhaustive +// expectation computed from the live vectors, and ghost uids must never appear. +// --------------------------------------------------------------------------- + +const ( + lifecyclePred = "lifecycle_vec" + lifecycleDim = 4 + lifecycleK = 10 +) + +func vecLiteral(v []float32) string { + parts := make([]string, len(v)) + for i, f := range v { + parts[i] = strconv.FormatFloat(float64(f), 'f', -1, 32) + } + return "[" + strings.Join(parts, ", ") + "]" +} + +// hnswExpected computes the exhaustive euclidean top-k over the live vectors +// with the server's exact orientation (score = 1/(1+distance)), ordered the +// way the query pipeline emits results: top-k selected by (score desc, uid +// asc), then sorted uid-ascending. +func hnswExpected(live map[uint64][]float32, qv []float32, k int) []scoredUID { + all := make([]scoredUID, 0, len(live)) + for uid, vec := range live { + var sum float64 + for i := range vec { + d := float64(vec[i]) - float64(qv[i]) + sum += d * d + } + all = append(all, scoredUID{uid: uid, score: 1.0 / (1.0 + math.Sqrt(sum))}) + } + sort.Slice(all, func(i, j int) bool { + if all[i].score != all[j].score { + return all[i].score > all[j].score + } + return all[i].uid < all[j].uid + }) + if k < len(all) { + all = all[:k] + } + sort.Slice(all, func(i, j int) bool { return all[i].uid < all[j].uid }) + return all +} + +func fetchSimilarTo(dg *dgraphapi.GrpcClient, k int, qv []float32) (map[uint64]float64, error) { + q := fmt.Sprintf(`{ + v as var(func: similar_to(%s, %d, "%s")) + q(func: uid(v)) { uid score: val(v) } + }`, lifecyclePred, k, vecLiteral(qv)) + resp, err := dg.Query(q) + if err != nil { + return nil, err + } + return parseScored(resp.GetJson()) +} + +// assertLifecycleState asserts, against the live-vector map: +// - exact top-k equality (uids and scores) with the exhaustive expectation, +// - a wide query (k > corpus size) returns only live uids (no ghosts), each +// with its exhaustive score, and never any uid from `ghosts`. +func assertLifecycleState(ct *assert.CollectT, dg *dgraphapi.GrpcClient, + live map[uint64][]float32, qv []float32, ghosts []uint64) { + + want := hnswExpected(live, qv, lifecycleK) + got, err := fetchSimilarTo(dg, lifecycleK, qv) + if !assert.NoError(ct, err) { + return + } + if assert.Len(ct, got, len(want), "top-k result count mismatch") { + for _, w := range want { + g, ok := got[w.uid] + if assert.True(ct, ok, "expected uid 0x%x in top-%d", w.uid, lifecycleK) { + assert.InDelta(ct, w.score, g, 1e-6, "score mismatch for uid 0x%x", w.uid) + } + } + } + + wide, err := fetchSimilarTo(dg, len(live)+30, qv) + if !assert.NoError(ct, err) { + return + } + exhaustive := hnswExpected(live, qv, len(live)) + exhaustiveByUID := make(map[uint64]float64, len(exhaustive)) + for _, e := range exhaustive { + exhaustiveByUID[e.uid] = e.score + } + for uid, score := range wide { + wantScore, isLive := exhaustiveByUID[uid] + if assert.True(ct, isLive, "ghost uid 0x%x returned by similar_to (not a live vector)", uid) { + assert.InDelta(ct, wantScore, score, 1e-6, "stale score for uid 0x%x", uid) + } + } + for _, ghost := range ghosts { + _, present := wide[ghost] + assert.False(ct, present, "deleted uid 0x%x reappeared in similar_to results", ghost) + } + // The exhaustive winner must always be retrievable. + if len(exhaustive) > 0 { + best := exhaustive[0] + for _, e := range exhaustive { + if e.score > best.score { + best = e + } + } + _, ok := wide[best.uid] + assert.True(ct, ok, "exhaustive winner uid 0x%x missing from wide similar_to", best.uid) + } +} + +func TestHNSWLifecycleParity(t *testing.T) { + conf := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1).WithReplicas(1) + c, err := dgraphtest.NewLocalCluster(conf) + require.NoError(t, err) + defer func() { c.Cleanup(t.Failed()) }() + require.NoError(t, c.Start()) + + dg, cleanup, err := c.Client() + require.NoError(t, err) + defer cleanup() + + require.NoError(t, dg.DropAll()) + require.NoError(t, dg.SetupSchema(fmt.Sprintf( + `%s: float32vector @index(hnsw(metric:"euclidean")) .`, lifecyclePred))) + require.NoError(t, c.AssignUids(dg.Dgraph, 128)) + + qv := []float32{0, 0, 0, 0} + live := make(map[uint64][]float32) + + setVec := func(uid uint64, v []float32) { + _, err := dg.Mutate(&api.Mutation{ + SetNquads: []byte(fmt.Sprintf("<0x%x> <%s> %q .", uid, lifecyclePred, vecLiteral(v))), + CommitNow: true, + }) + require.NoError(t, err) + live[uid] = v + } + delVec := func(uid uint64) { + _, err := dg.Mutate(&api.Mutation{ + DelNquads: []byte(fmt.Sprintf("<0x%x> <%s> * .", uid, lifecyclePred)), + CommitNow: true, + }) + require.NoError(t, err) + delete(live, uid) + } + verify := func(label string, ghosts ...uint64) { + t.Helper() + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assertLifecycleState(ct, dg, live, qv, ghosts) + }, 30*time.Second, 500*time.Millisecond, "step %q did not converge", label) + } + + // Step 0: initial corpus — 30 vectors on a line; uid 1 (distance 3) wins. + var initial []string + for i := 1; i <= 30; i++ { + v := []float32{float32(3 * i), 0, 0, 0} + live[uint64(i)] = v + initial = append(initial, + fmt.Sprintf("<0x%x> <%s> %q .", uint64(i), lifecyclePred, vecLiteral(v))) + } + _, err = dg.Mutate(&api.Mutation{ + SetNquads: []byte(strings.Join(initial, "\n")), CommitNow: true, + }) + require.NoError(t, err) + verify("initial") + require.Equal(t, uint64(1), hnswExpected(live, qv, 1)[0].uid) + + // Step 1: upsert uid 5 (was distance 15) to distance 1 — new expected winner. + setVec(5, []float32{-1, 0, 0, 0}) + require.Equal(t, uint64(5), hnswExpected(live, qv, 1)[0].uid) + verify("upsert-winner-flip") + + // Step 2: delete the winner. uid 1 must win again; uid 5 must be a ghost never returned. + delVec(5) + require.Equal(t, uint64(1), hnswExpected(live, qv, 1)[0].uid) + verify("delete-winner", 5) + + // Step 3: reinsert the SAME uid with a DIFFERENT vector (distance 2, winner + // again). Its score must reflect the new vector, not the pre-delete one. + setVec(5, []float32{-2, 0, 0, 0}) + require.Equal(t, uint64(5), hnswExpected(live, qv, 1)[0].uid) + verify("reinsert-same-uid") + + // Step 4: kill -9 the alpha, restart, and require convergence to the exact + // same exhaustive expectation with no ghosts. + require.NoError(t, c.KillAlpha(0)) + require.NoError(t, c.StartAlpha(0)) + require.NoError(t, c.HealthCheck(false)) + + dg2, cleanup2, err := c.Client() + require.NoError(t, err) + defer cleanup2() + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assertLifecycleState(ct, dg2, live, qv, []uint64{}) + }, 90*time.Second, 2*time.Second, "post-restart state did not converge") +} diff --git a/systest/bm25lifecycle/rebuild_crash_test.go b/systest/bm25lifecycle/rebuild_crash_test.go new file mode 100644 index 00000000000..2cdf9014b14 --- /dev/null +++ b/systest/bm25lifecycle/rebuild_crash_test.go @@ -0,0 +1,537 @@ +//go:build integration2 + +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package main + +// Wave-2 lifecycle tests for BM25 index rebuilds (questions Q16 and Q20-lite). +// +// Q16 (KNOWN-FAILING when the overlap manifests): a background @index(bm25) +// rebuild DropPrefix'es the stats buckets and flushes its absolute totals at the +// rebuild's startTs. Live mutations committing between the DropPrefix and the +// flush read an EMPTY bucket, and their zero-based read-modify-write absolute +// total shadows the rebuild's flush via MVCC (higher commitTs wins). The final +// corpus stats therefore UNDERCOUNT, which shows up as a wrong IDF in every +// BM25 score. We assert final scores against an independently computed +// closed-form expectation, which discriminates regardless of mechanism. +// +// Q20-lite (expected PASS): the forever-empty hazard was refuted — the schema +// commits only after BuildIndexes succeeds, and Raft replay re-runs the rebuild +// after a crash. This is a crash-RECOVERY sanity test: SIGKILL the alpha +// mid-rebuild, restart it, and poll until bm25 queries return the full expected +// result set with exact closed-form scores. +// +// Corpus design: every document tokenizes to exactly 2 terms with tf=1 each +// (distinct non-stopword, stem-stable nonsense words). Then avgDL == dl == 2, +// so the BM25 term factor (k+1)*tf/(k*(1-b+b*dl/avgDL)+tf) collapses to 1 and +// every score equals the smoothed IDF: ln(1 + (N-df+0.5)/(df+0.5)). docCount +// (via IDF) and totalTerms (via avgDL) are thus both pinned by score equality, +// even though ReadBM25Stats is not client-reachable. + +import ( + "context" + "encoding/json" + "fmt" + "io" + "math" + "net/http" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/dgraph-io/dgo/v250/protos/api" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/dgraphapi" + "github.com/dgraph-io/dgraph/v25/dgraphtest" +) + +type scoredNode struct { + UID string `json:"uid"` + Score float64 `json:"score"` +} + +// bm25Probe runs an unpaginated bm25() root function binding the score to a +// value variable, returning every match with its score. No `first` is used so +// the worker takes the score-all path (topK == 0) and emits the full corpus of +// matches. +func bm25Probe(dg *dgraphapi.GrpcClient, pred, term string) ([]scoredNode, error) { + q := fmt.Sprintf(`{ + s as var(func: bm25(%s, %q)) + q(func: uid(s)) { + uid + score: val(s) + } + }`, pred, term) + resp, err := dg.Query(q) + if err != nil { + return nil, err + } + var out struct { + Q []scoredNode `json:"q"` + } + if err := json.Unmarshal(resp.GetJson(), &out); err != nil { + return nil, err + } + return out.Q, nil +} + +// expectedBM25Score mirrors worker/bm25wand.go exactly: smoothed IDF +// (math.Log1p((N-df+0.5)/(df+0.5))) times the BM25 term factor with default +// k=1.2, b=0.75, and our fixed dl=2, tf=1. With totalTerms/nDocs == 2 the term +// factor is 1 and the score is the IDF alone. +func expectedBM25Score(nDocs, df, totalTerms float64) float64 { + const k, b = 1.2, 0.75 + const dl, tf = 2.0, 1.0 + avgDL := totalTerms / nDocs + idf := math.Log1p((nDocs - df + 0.5) / (df + 0.5)) + return idf * (k + 1) * tf / (k*(1-b+b*dl/avgDL) + tf) +} + +// loadTwoTokenDocs loads `count` docs all carrying the same two-token text, +// in batches of 1000 (the predicate has no bm25 index at load time, so batch +// size cannot trip the Q15/Q19 intra-proposal stats races). Returns the +// assigned uids in blank-node order. +func loadTwoTokenDocs(dg *dgraphapi.GrpcClient, pred, text, labelPrefix string, count int) ([]string, error) { + const batchSize = 1000 + uids := make([]string, 0, count) + for start := 0; start < count; start += batchSize { + end := start + batchSize + if end > count { + end = count + } + var sb strings.Builder + for i := start; i < end; i++ { + fmt.Fprintf(&sb, "_:%s%d <%s> %q .\n", labelPrefix, i, pred, text) + } + resp, err := dg.Mutate(&api.Mutation{SetNquads: []byte(sb.String()), CommitNow: true}) + if err != nil { + return nil, err + } + for i := start; i < end; i++ { + uid, ok := resp.Uids[fmt.Sprintf("%s%d", labelPrefix, i)] + if !ok { + return nil, fmt.Errorf("no uid assigned for blank node %s%d", labelPrefix, i) + } + uids = append(uids, uid) + } + } + return uids, nil +} + +// mutateWithRetry commits a single mutation, retrying transaction aborts. +// Single-doc txns on a bm25 predicate conflict on the shared uid%32 stats +// buckets, so aborts are expected under concurrency and must be retried. +func mutateWithRetry(dg *dgraphapi.GrpcClient, mu *api.Mutation) (*api.Response, error) { + var lastErr error + for attempt := 0; attempt < 100; attempt++ { + resp, err := dg.Mutate(mu) + if err == nil { + return resp, nil + } + lastErr = err + low := strings.ToLower(err.Error()) + if !strings.Contains(low, "abort") && !strings.Contains(low, "retry") && + !strings.Contains(low, "conflict") { + return nil, err + } + time.Sleep(50 * time.Millisecond) // retry backoff, not a convergence wait + } + return nil, lastErr +} + +// alphaIndexingInProgress reports whether alpha 0's /health?all lists `pred` +// as currently being indexed (schema.GetIndexingPredicates). Best-effort: +// any error reads as false. +func alphaIndexingInProgress(c *dgraphtest.LocalCluster, pred string) bool { + port, err := c.GetAlphaHttpPublicPort(0) + if err != nil { + return false + } + httpc := http.Client{Timeout: 5 * time.Second} + resp, err := httpc.Get("http://0.0.0.0:" + port + "/health?all") + if err != nil { + return false + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(resp.Body) + if err != nil { + return false + } + return strings.Contains(string(body), `"indexing"`) && strings.Contains(string(body), pred) +} + +// pollUntil runs f every tick until it returns true or the timeout elapses. +// Used for best-effort window detection (NOT for correctness assertions — +// those use require/assert.Eventually). +func pollUntil(timeout, tick time.Duration, f func() bool) bool { + deadline := time.Now().Add(timeout) + ticker := time.NewTicker(tick) + defer ticker.Stop() + for { + if f() { + return true + } + if time.Now().After(deadline) { + return false + } + <-ticker.C + } +} + +// TestBM25RebuildUnderLiveWrites — Q16, KNOWN-FAILING when the rebuild/live- +// write overlap manifests. +// +// Load 5000 unindexed docs, add @index(bm25) with RunInBackground, and while +// the rebuild runs commit concurrent adds/updates/deletes in single-doc txns. +// Afterwards every score must equal the closed-form expectation computed from +// the true final corpus (see the ground-truth math below). The failing +// signature is a score whose implied docCount collapsed toward the number of +// live writes (bucket totals clobbered by zero-based RMW totals shadowing the +// rebuild's absolute flush), or a mid-rebuild-committed doc that bm25() never +// finds. +func TestBM25RebuildUnderLiveWrites(t *testing.T) { + // KNOWN-FAILING (Q16), gated: live writes committing between the rebuild's + // DropPrefix and its absolute stats flush read an EMPTY bucket and their + // zero-based RMW total shadows the flush via MVCC — empirically confirmed + // (uid 0x5b scored 0.0062, closed-form expectation 4.137; docCount collapsed). + // The fix needs a Zero-mediated conflicting flush or delta-merged stats; this + // test is its acceptance gate. Set BM25_KNOWN_FAILING=1 to run. + if os.Getenv("BM25_KNOWN_FAILING") == "" { + t.Skip("KNOWN-FAILING (Q16 rebuild-vs-live stats undercount); set BM25_KNOWN_FAILING=1 to run") + } + const ( + pred = "rebuild_text" + baseDocs = 5000 + liveAdds = 60 + liveUpdates = 20 + liveDeletes = 20 + ) + // Final ground truth: + // docCount = 5000 + 60 adds - 20 deletes = 5040 + // totalTerms = 5000*2 + 60*2 - 20*2 (deletes; updates net 0) = 10080 + // df(quokka) = 60 adds + 20 updates = 80 + // df(zeppelin) = 5000 - 20 deleted - 20 updated = 4960 + const ( + wantDocCount = float64(baseDocs + liveAdds - liveDeletes) + wantTotalTerms = float64(2 * (baseDocs + liveAdds - liveDeletes)) + wantQuokkaDF = float64(liveAdds + liveUpdates) + wantZeppelinDF = float64(baseDocs - liveDeletes - liveUpdates) + ) + + conf := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1).WithReplicas(1) + c, err := dgraphtest.NewLocalCluster(conf) + require.NoError(t, err) + defer func() { c.Cleanup(t.Failed()) }() + require.NoError(t, c.Start()) + + dg, cleanup, err := c.Client() + require.NoError(t, err) + defer cleanup() + + require.NoError(t, dg.DropAll()) + require.NoError(t, dg.SetupSchema(pred+`: string .`)) + + baseUids, err := loadTwoTokenDocs(dg, pred, "zeppelin marzipan", "b", baseDocs) + require.NoError(t, err) + require.Len(t, baseUids, baseDocs) + + // Sanity: all base docs are present before the index exists. + resp, err := dg.Query(fmt.Sprintf(`{ q(func: has(%s)) { count(uid) } }`, pred)) + require.NoError(t, err) + require.Contains(t, string(resp.GetJson()), fmt.Sprintf(`"count":%d`, baseDocs)) + + // Kick off the background rebuild. The alter is acknowledged before + // BuildIndexes runs, so the live writes below overlap the rebuild window. + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + require.NoError(t, dg.Alter(ctx, &api.Operation{ + Schema: pred + `: string @index(bm25) .`, + RunInBackground: true, + })) + + // Observability only: record whether /health?all reported the predicate as + // indexing while the live writes were in flight, so a PASS can be + // distinguished from "rebuild finished before the writes started". + var sawIndexing atomic.Bool + stopMonitor := make(chan struct{}) + monitorDone := make(chan struct{}) + go func() { + defer close(monitorDone) + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-stopMonitor: + return + case <-ticker.C: + if alphaIndexingInProgress(c, pred) { + sawIndexing.Store(true) + } + } + } + }() + + // Live writes, all single-doc txns (< 512 edges, so the Q19 intra-proposal + // batching race cannot confound this test), three concurrent streams. + updateUids := baseUids[:liveUpdates] + deleteUids := baseUids[liveUpdates : liveUpdates+liveDeletes] + addUids := make([]string, liveAdds) + errs := make(chan error, 3) + var wg sync.WaitGroup + wg.Add(3) + go func() { // adds: new docs matching "quokka" + defer wg.Done() + for i := 0; i < liveAdds; i++ { + r, err := mutateWithRetry(dg, &api.Mutation{ + SetNquads: []byte(fmt.Sprintf("_:n <%s> \"quokka lutefisk\" .", pred)), + CommitNow: true, + }) + if err != nil { + errs <- fmt.Errorf("live add %d: %w", i, err) + return + } + addUids[i] = r.Uids["n"] + } + errs <- nil + }() + go func() { // updates: flip existing docs from "zeppelin marzipan" to "quokka gruyere" + defer wg.Done() + for i, uid := range updateUids { + _, err := mutateWithRetry(dg, &api.Mutation{ + SetNquads: []byte(fmt.Sprintf("<%s> <%s> \"quokka gruyere\" .", uid, pred)), + CommitNow: true, + }) + if err != nil { + errs <- fmt.Errorf("live update %d: %w", i, err) + return + } + } + errs <- nil + }() + go func() { // deletes: remove existing docs entirely + defer wg.Done() + for i, uid := range deleteUids { + _, err := mutateWithRetry(dg, &api.Mutation{ + DelNquads: []byte(fmt.Sprintf("<%s> <%s> * .", uid, pred)), + CommitNow: true, + }) + if err != nil { + errs <- fmt.Errorf("live delete %d: %w", i, err) + return + } + } + errs <- nil + }() + wg.Wait() + close(stopMonitor) + <-monitorDone + for i := 0; i < 3; i++ { + require.NoError(t, <-errs) + } + t.Logf("Q16: indexing observed during live writes: %v "+ + "(false means the rebuild won the race and the undercount window never opened)", + sawIndexing.Load()) + + // Wait for the rebuild to finish: the bm25() query errors with "not + // indexed" until the schema commits after BuildIndexes. Converge on the + // full expected match count, then assert scores separately for a crisp + // failure signature. + lastState := "" + converged := assert.Eventually(t, func() bool { + nodes, err := bm25Probe(dg, pred, "quokka") + state := "" + if err != nil { + state = "query error: " + err.Error() + } else { + state = fmt.Sprintf("%d quokka matches (want %d)", len(nodes), liveAdds+liveUpdates) + } + if state != lastState { + t.Logf("Q16 convergence: %s", state) + lastState = state + } + return err == nil && len(nodes) == liveAdds+liveUpdates + }, 90*time.Second, 500*time.Millisecond) + require.True(t, converged, + "Q16: bm25 never returned the full live-write result set; last state: %s "+ + "(a mid-rebuild-committed doc that bm25() never finds is a Q16 failure signature)", + lastState) + + // Final assertions against the independently computed expectation. + quokkaNodes, err := bm25Probe(dg, pred, "quokka") + require.NoError(t, err) + require.Len(t, quokkaNodes, liveAdds+liveUpdates) + + expectedQuokkaUids := make(map[string]struct{}, liveAdds+liveUpdates) + for _, uid := range addUids { + expectedQuokkaUids[uid] = struct{}{} + } + for _, uid := range updateUids { + expectedQuokkaUids[uid] = struct{}{} + } + for _, n := range quokkaNodes { + _, ok := expectedQuokkaUids[n.UID] + require.True(t, ok, "unexpected uid %s in quokka results (deleted or phantom doc)", n.UID) + } + + wantQuokkaScore := expectedBM25Score(wantDocCount, wantQuokkaDF, wantTotalTerms) + for _, n := range quokkaNodes { + // KNOWN-FAILING (Q16): when live writes overlap the rebuild's + // DropPrefix→flush window, the stats undercount makes the observed + // score imply a docCount collapsed toward the number of live writes. + require.InDeltaf(t, wantQuokkaScore, n.Score, 1e-6, + "Q16 UNDERCOUNT SIGNATURE: uid %s scored %v, want %v "+ + "(closed-form with docCount=%v totalTerms=%v df=%v; a lower observed score implies "+ + "the rebuild's absolute stats flush was shadowed by a zero-based live-write RMW)", + n.UID, n.Score, wantQuokkaScore, wantDocCount, wantTotalTerms, wantQuokkaDF) + } + + // Cross-check via the other term: count pins deletes/updates having taken + // effect in the term postings, score re-pins the same (N, avgDL). + zeppelinNodes, err := bm25Probe(dg, pred, "zeppelin") + require.NoError(t, err) + require.Len(t, zeppelinNodes, int(wantZeppelinDF), + "zeppelin match count wrong: deleted/updated docs must not match after rebuild") + wantZeppelinScore := expectedBM25Score(wantDocCount, wantZeppelinDF, wantTotalTerms) + for _, n := range zeppelinNodes { + require.InDeltaf(t, wantZeppelinScore, n.Score, 1e-6, + "Q16 UNDERCOUNT SIGNATURE (zeppelin channel): uid %s scored %v, want %v", + n.UID, n.Score, wantZeppelinScore) + } +} + +// TestBM25CrashRecoveryMidRebuild — Q20-lite, expected PASS. +// +// SIGKILL the sole alpha while a fresh @index(bm25) rebuild is running, restart +// it, and require convergence: the schema mutation lives in the Raft log, so +// replay re-runs the rebuild and bm25() must eventually return the full +// expected result set with exact closed-form scores. (The forever-empty hazard +// was refuted: the schema commits only after BuildIndexes succeeds.) +func TestBM25CrashRecoveryMidRebuild(t *testing.T) { + const ( + pred = "crash_text" + fillDocs = 2950 // "zeppelin marzipan" + probeDocs = 50 // "quokka lutefisk" + ) + const ( + wantDocCount = float64(fillDocs + probeDocs) // 3000 + wantTotalTerms = float64(2 * (fillDocs + probeDocs)) + wantQuokkaDF = float64(probeDocs) + wantZeppelinDF = float64(fillDocs) + ) + + conf := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1).WithReplicas(1) + c, err := dgraphtest.NewLocalCluster(conf) + require.NoError(t, err) + defer func() { c.Cleanup(t.Failed()) }() + require.NoError(t, c.Start()) + + dg, cleanup, err := c.Client() + require.NoError(t, err) + defer cleanup() + + require.NoError(t, dg.DropAll()) + require.NoError(t, dg.SetupSchema(pred+`: string .`)) + + _, err = loadTwoTokenDocs(dg, pred, "zeppelin marzipan", "f", fillDocs) + require.NoError(t, err) + probeUids, err := loadTwoTokenDocs(dg, pred, "quokka lutefisk", "p", probeDocs) + require.NoError(t, err) + require.Len(t, probeUids, probeDocs) + + // Start the rebuild in the background; the ack precedes BuildIndexes. + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + require.NoError(t, dg.Alter(ctx, &api.Operation{ + Schema: pred + `: string @index(bm25) .`, + RunInBackground: true, + })) + + // Best-effort: catch the rebuild in flight before killing. If the rebuild + // outruns us the test degrades to a plain crash-recovery check, which is + // still a valid (weaker) run — log which one we got. + caughtMidRebuild := pollUntil(15*time.Second, 20*time.Millisecond, func() bool { + return alphaIndexingInProgress(c, pred) + }) + t.Logf("Q20: SIGKILL with indexing in progress: %v", caughtMidRebuild) + + require.NoError(t, c.KillAlpha(0)) + require.NoError(t, c.StartAlpha(0)) + require.NoError(t, c.HealthCheck(false)) + + // The pre-kill client's connection may be stale; dial fresh. + dg2, cleanup2, err := c.Client() + require.NoError(t, err) + defer cleanup2() + + // Converge: Raft replay must re-run the rebuild; poll until bm25 returns + // the complete probe set with exact scores. Generous timeout, no fixed + // sleeps. + wantQuokkaScore := expectedBM25Score(wantDocCount, wantQuokkaDF, wantTotalTerms) + lastState := "" + converged := assert.Eventually(t, func() bool { + nodes, err := bm25Probe(dg2, pred, "quokka") + state := "" + switch { + case err != nil: + state = "query error: " + err.Error() + case len(nodes) != probeDocs: + state = fmt.Sprintf("%d quokka matches (want %d)", len(nodes), probeDocs) + default: + maxDiff := 0.0 + for _, n := range nodes { + if d := math.Abs(n.Score - wantQuokkaScore); d > maxDiff { + maxDiff = d + } + } + if maxDiff > 1e-6 { + state = fmt.Sprintf("scores off by up to %v (want %v)", maxDiff, wantQuokkaScore) + } else { + state = "converged" + } + } + if state != lastState { + t.Logf("Q20 convergence: %s", state) + lastState = state + } + return state == "converged" + }, 120*time.Second, time.Second) + require.True(t, converged, + "Q20: bm25 never converged to the full exact result set after crash-restart; last state: %s "+ + "(schema listing bm25 while results stay empty/short would be the lost-rebuild signature)", + lastState) + + // Final exact assertions for crisp failure messages. + nodes, err := bm25Probe(dg2, pred, "quokka") + require.NoError(t, err) + require.Len(t, nodes, probeDocs) + expectedUids := make(map[string]struct{}, probeDocs) + for _, uid := range probeUids { + expectedUids[uid] = struct{}{} + } + for _, n := range nodes { + _, ok := expectedUids[n.UID] + require.True(t, ok, "unexpected uid %s in post-recovery quokka results", n.UID) + require.InDeltaf(t, wantQuokkaScore, n.Score, 1e-6, + "post-recovery score for uid %s is %v, want closed-form %v (stats not exact after replayed rebuild)", + n.UID, n.Score, wantQuokkaScore) + } + + // Cross-check the other term: full corpus stats must be exact, not merely + // the probe docs' postings. + zeppelinNodes, err := bm25Probe(dg2, pred, "zeppelin") + require.NoError(t, err) + require.Len(t, zeppelinNodes, fillDocs) + wantZeppelinScore := expectedBM25Score(wantDocCount, wantZeppelinDF, wantTotalTerms) + for _, n := range zeppelinNodes { + require.InDeltaf(t, wantZeppelinScore, n.Score, 1e-6, + "post-recovery zeppelin score for uid %s is %v, want %v", n.UID, n.Score, wantZeppelinScore) + } +} diff --git a/systest/bm25lifecycle/smoke_test.go b/systest/bm25lifecycle/smoke_test.go new file mode 100644 index 00000000000..91bf3954047 --- /dev/null +++ b/systest/bm25lifecycle/smoke_test.go @@ -0,0 +1,47 @@ +//go:build integration2 + +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package main + +import ( + "testing" + + "github.com/dgraph-io/dgo/v250/protos/api" + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/dgraphtest" +) + +// TestBM25LifecycleSmoke proves the LocalCluster harness works for this package: +// single group, bm25 index, one query round-trip. +func TestBM25LifecycleSmoke(t *testing.T) { + conf := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1).WithReplicas(1) + c, err := dgraphtest.NewLocalCluster(conf) + require.NoError(t, err) + defer func() { c.Cleanup(t.Failed()) }() + require.NoError(t, c.Start()) + + dg, cleanup, err := c.Client() + require.NoError(t, err) + defer cleanup() + + require.NoError(t, dg.DropAll()) + require.NoError(t, dg.SetupSchema(`smoke_text: string @index(bm25) .`)) + + _, err = dg.Mutate(&api.Mutation{ + SetNquads: []byte(` + _:a "zeppelin zeppelin" . + _:b "zeppelin marzipan" . + `), + CommitNow: true, + }) + require.NoError(t, err) + + resp, err := dg.Query(`{ q(func: bm25(smoke_text, "zeppelin")) { uid } }`) + require.NoError(t, err) + require.Contains(t, string(resp.GetJson()), "uid") +} diff --git a/systest/integration2/bulk_loader_test.go b/systest/integration2/bulk_loader_test.go index 48b59b01d8a..bd005476d02 100644 --- a/systest/integration2/bulk_loader_test.go +++ b/systest/integration2/bulk_loader_test.go @@ -10,6 +10,7 @@ package main import ( "os" "path/filepath" + "strings" "testing" "time" @@ -122,6 +123,62 @@ func TestBulkLoaderSkipReducePhase(t *testing.T) { }`, string(data))) } +// TestBulkLoaderBM25 verifies the BM25 index is correctly built by the bulk loader: +// term postings must carry their packed (term frequency, document length) value, and +// the corpus statistics must be written, or bm25() queries return nothing. It loads a +// small corpus via bulk, then checks that all documents containing the term are found +// and that the densest/shortest document ranks first. +func TestBulkLoaderBM25(t *testing.T) { + conf := dgraphtest.NewClusterConfig().WithNumAlphas(1).WithNumZeros(1). + WithACL(time.Hour).WithReplicas(1).WithBulkLoadOutDir(t.TempDir()) + c, err := dgraphtest.NewLocalCluster(conf) + require.NoError(t, err) + defer func() { c.Cleanup(t.Failed()) }() + + require.NoError(t, c.StartZero(0)) + require.NoError(t, c.HealthCheck(true)) + + baseDir := t.TempDir() + schemaFile := filepath.Join(baseDir, "bm25.schema") + require.NoError(t, os.WriteFile(schemaFile, + []byte("description_bm25: string @index(bm25) .\n"), os.ModePerm)) + + dataFile := filepath.Join(baseDir, "bm25.rdf") + rdf := ` + <0x1> "the quick brown fox jumps over the lazy dog" . + <0x2> "fox fox fox" . + <0x3> "the lazy dog sleeps in the warm sun all day" . + <0x4> "quick brown foxes are agile animals" . + ` + require.NoError(t, os.WriteFile(dataFile, []byte(rdf), os.ModePerm)) + + require.NoError(t, c.BulkLoad(dgraphtest.BulkOpts{ + DataFiles: []string{dataFile}, + SchemaFiles: []string{schemaFile}, + })) + + require.NoError(t, c.Start()) + + hc, err := c.HTTPClient() + require.NoError(t, err) + require.NoError(t, hc.LoginIntoNamespace(dgraphapi.DefaultUser, + dgraphapi.DefaultPassword, x.RootNamespace)) + + // Both documents containing "fox" (0x1 and 0x2) must be found — proves the term + // postings and corpus statistics survived the bulk build. + data, err := hc.PostDqlQuery(`{ q(func: bm25(description_bm25, "fox")) { count(uid) } }`) + require.NoError(t, err) + require.Contains(t, string(data), `"count":3`, + "bulk-loaded bm25 index must find every document containing the term") + + // The all-"fox" document (0x2, tf=3, shortest) must rank first. + data, err = hc.PostDqlQuery( + `{ q(func: bm25(description_bm25, "fox"), first: 1) { description_bm25 } }`) + require.NoError(t, err) + require.True(t, strings.Contains(string(data), "fox fox fox"), + "densest, shortest document must rank first after bulk load; got: %s", string(data)) +} + func TestBulkLoaderNoDqlSchema(t *testing.T) { conf := dgraphtest.NewClusterConfig().WithNumAlphas(2).WithNumZeros(1). WithACL(time.Hour).WithReplicas(1).WithBulkLoadOutDir(t.TempDir()) diff --git a/tok/hnsw/helper.go b/tok/hnsw/helper.go index 39d72d8f5e7..890ca170480 100644 --- a/tok/hnsw/helper.go +++ b/tok/hnsw/helper.go @@ -214,6 +214,24 @@ type SimilarityType[T c.Float] struct { isSimilarityMetric bool } +// similarityScore converts a heap element's metric-domain value into a +// higher-is-better similarity score suitable for ranking and score fusion. +// +// - Cosine / dot product (isSimilarityMetric): the value is already a similarity +// where higher is better, so it is returned as-is. +// - Euclidean: the value is a squared L2 distance where lower is better. It is +// mapped to 1/(1+d) in (0,1], which is monotonically decreasing in distance, so +// higher is better and the result is well-behaved under linear normalization. +// +// Keeping this orientation in one place lets every caller (and hybrid-search +// fusion) treat vector scores with the same higher-is-better convention as BM25. +func (s SimilarityType[T]) similarityScore(value T) float64 { + if s.isSimilarityMetric { + return float64(value) + } + return 1.0 / (1.0 + float64(value)) +} + func GetSimType[T c.Float](indexType string, floatBits int) SimilarityType[T] { switch { case indexType == Euclidean: diff --git a/tok/hnsw/persistent_hnsw.go b/tok/hnsw/persistent_hnsw.go index 5658800e579..e7b8561e24e 100644 --- a/tok/hnsw/persistent_hnsw.go +++ b/tok/hnsw/persistent_hnsw.go @@ -266,6 +266,20 @@ func (ph *persistentHNSW[T]) SearchWithOptions( maxResults int, opts index.VectorIndexOptions[T], ) ([]uint64, error) { + uids, _, err := ph.SearchWithOptionsScored(ctx, c, query, maxResults, opts) + return uids, err +} + +// SearchWithOptionsScored is SearchWithOptions that also returns a higher-is-better +// similarity score for each returned uid (positionally aligned). See +// index.ScoredSearchOptions. +func (ph *persistentHNSW[T]) SearchWithOptionsScored( + ctx context.Context, + c index.CacheType, + query []T, + maxResults int, + opts index.VectorIndexOptions[T], +) ([]uint64, []float64, error) { if opts.Filter == nil { opts.Filter = index.AcceptAll[T] } @@ -279,7 +293,7 @@ func (ph *persistentHNSW[T]) SearchWithOptions( var startVec []T entry, err := ph.PickStartNode(ctx, c, &startVec) if err != nil { - return nil, err + return nil, nil, err } // Upper layers use efUpper (override if provided) @@ -296,13 +310,13 @@ func (ph *persistentHNSW[T]) SearchWithOptions( layerResult, err := ph.searchPersistentLayer( c, level, entry, startVec, query, filterOut, efUpper, opts.Filter) if err != nil { - return nil, err + return nil, nil, err } layerResult.updateFinalMetrics(r) entry = layerResult.bestNeighbor().index layerResult.updateFinalPath(r) if err = ph.getVecFromUid(entry, c, &startVec); err != nil { - return nil, err + return nil, nil, err } } @@ -315,13 +329,14 @@ func (ph *persistentHNSW[T]) SearchWithOptions( layerResult, err := ph.searchPersistentLayer( c, ph.maxLevels-1, entry, startVec, query, filterOut, candidateK, opts.Filter) if err != nil { - return nil, err + return nil, nil, err } layerResult.updateFinalMetrics(r) layerResult.updateFinalPath(r) // Build final neighbor list with optional threshold, limited to maxResults. res := make([]uint64, 0, maxResults) + scores := make([]float64, 0, maxResults) for _, n := range layerResult.neighbors { if maxResults == 0 { break @@ -347,23 +362,38 @@ func (ph *persistentHNSW[T]) SearchWithOptions( } } res = append(res, n.index) + scores = append(scores, ph.simType.similarityScore(n.value)) if len(res) >= maxResults { break } } r.Metrics[searchTime] = uint64(time.Now().UnixMilli() - start) - return res, nil + return res, scores, nil } // SearchWithUidAndOptions is analogous to SearchWithUid but applies per‑call options. func (ph *persistentHNSW[T]) SearchWithUidAndOptions( - _ context.Context, + ctx context.Context, c index.CacheType, queryUid uint64, maxResults int, opts index.VectorIndexOptions[T], ) ([]uint64, error) { + uids, _, err := ph.SearchWithUidAndOptionsScored(ctx, c, queryUid, maxResults, opts) + return uids, err +} + +// SearchWithUidAndOptionsScored is SearchWithUidAndOptions that also returns a +// higher-is-better similarity score for each returned uid (positionally aligned). +// See index.ScoredSearchOptions. +func (ph *persistentHNSW[T]) SearchWithUidAndOptionsScored( + _ context.Context, + c index.CacheType, + queryUid uint64, + maxResults int, + opts index.VectorIndexOptions[T], +) ([]uint64, []float64, error) { if opts.Filter == nil { opts.Filter = index.AcceptAll[T] } @@ -373,12 +403,12 @@ func (ph *persistentHNSW[T]) SearchWithUidAndOptions( var queryVec []T if err := ph.getVecFromUid(queryUid, c, &queryVec); err != nil { if errors.Is(err, errFetchingPostingList) { - return []uint64{}, nil + return []uint64{}, []float64{}, nil } - return []uint64{}, err + return []uint64{}, []float64{}, err } if len(queryVec) == 0 { - return []uint64{}, nil + return []uint64{}, []float64{}, nil } filterOut := !opts.Filter(queryVec, queryVec, queryUid) candidateK := maxResults @@ -388,9 +418,10 @@ func (ph *persistentHNSW[T]) SearchWithUidAndOptions( lr, err := ph.searchPersistentLayer( c, ph.maxLevels-1, queryUid, queryVec, queryVec, filterOut, candidateK, opts.Filter) if err != nil { - return []uint64{}, err + return []uint64{}, []float64{}, err } res := make([]uint64, 0, maxResults) + scores := make([]float64, 0, maxResults) for _, n := range lr.neighbors { if maxResults == 0 { break @@ -413,11 +444,12 @@ func (ph *persistentHNSW[T]) SearchWithUidAndOptions( } } res = append(res, n.index) + scores = append(scores, ph.simType.similarityScore(n.value)) if len(res) >= maxResults { break } } - return res, nil + return res, scores, nil } // SearchWithUid searches the HNSW graph for the nearest neighbors of the query UID @@ -548,13 +580,56 @@ func (ph *persistentHNSW[T]) SearchWithPath( } layerResult.updateFinalMetrics(r) layerResult.updateFinalPath(r) - layerResult.addFinalNeighbors(r) + layerResult.addFinalNeighbors(r, ph.simType) t := time.Now().UnixMilli() elapsed := t - start r.Metrics[searchTime] = uint64(elapsed) return r, nil } +// SearchScored is Search that also returns a higher-is-better similarity score for +// each returned uid (positionally aligned with the neighbor uids). It preserves the +// exact candidate-exploration behavior of Search (unlike the options-based path), +// so scoring an otherwise plain query does not change which neighbors are returned. +func (ph *persistentHNSW[T]) SearchScored(ctx context.Context, c index.CacheType, query []T, + maxResults int, filter index.SearchFilter[T]) ([]uint64, []float64, error) { + r, err := ph.SearchWithPath(ctx, c, query, maxResults, filter) + if err != nil { + return nil, nil, err + } + return r.Neighbors, r.Distances, nil +} + +// SearchWithUidScored is SearchWithUid that also returns a higher-is-better +// similarity score for each returned uid (positionally aligned), preserving +// SearchWithUid's exact neighbor selection. +func (ph *persistentHNSW[T]) SearchWithUidScored(_ context.Context, c index.CacheType, + queryUid uint64, maxResults int, filter index.SearchFilter[T]) ([]uint64, []float64, error) { + var queryVec []T + if err := ph.getVecFromUid(queryUid, c, &queryVec); err != nil { + if errors.Is(err, errFetchingPostingList) { + return []uint64{}, []float64{}, nil + } + return []uint64{}, []float64{}, err + } + if len(queryVec) == 0 { + return []uint64{}, []float64{}, nil + } + shouldFilterOutQueryVec := !filter(queryVec, queryVec, queryUid) + r, err := ph.searchPersistentLayer( + c, ph.maxLevels-1, queryUid, queryVec, queryVec, shouldFilterOutQueryVec, maxResults, filter) + if err != nil { + return []uint64{}, []float64{}, err + } + uids := make([]uint64, 0, len(r.neighbors)) + scores := make([]float64, 0, len(r.neighbors)) + for _, n := range r.neighbors { + uids = append(uids, n.index) + scores = append(scores, ph.simType.similarityScore(n.value)) + } + return uids, scores, nil +} + // InsertToPersistentStorage inserts a node into the HNSW graph and returns the // traversal path and the edges created func (ph *persistentHNSW[T]) Insert(ctx context.Context, c index.CacheType, diff --git a/tok/hnsw/search_layer.go b/tok/hnsw/search_layer.go index 55140e7319a..81c529bd29b 100644 --- a/tok/hnsw/search_layer.go +++ b/tok/hnsw/search_layer.go @@ -113,10 +113,13 @@ func (slr *searchLayerResult[T]) updateFinalPath(r *index.SearchPathResult) { r.Path = append(r.Path, slr.path...) } -func (slr *searchLayerResult[T]) addFinalNeighbors(r *index.SearchPathResult) { +func (slr *searchLayerResult[T]) addFinalNeighbors(r *index.SearchPathResult, simType SimilarityType[T]) { for _, n := range slr.neighbors { if !n.filteredOut { r.Neighbors = append(r.Neighbors, n.index) + // Distances carries the higher-is-better similarity for each neighbor, + // positionally aligned with Neighbors, so scored searches can surface it. + r.Distances = append(r.Distances, simType.similarityScore(n.value)) } } } diff --git a/tok/hnsw/similarity_score_test.go b/tok/hnsw/similarity_score_test.go new file mode 100644 index 00000000000..481c97371b7 --- /dev/null +++ b/tok/hnsw/similarity_score_test.go @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package hnsw + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// TestSimilarityScoreOrientation verifies that every metric surfaces a +// higher-is-better similarity score, which native hybrid-search fusion relies on. +func TestSimilarityScoreOrientation(t *testing.T) { + cosine := GetSimType[float32](Cosine, 32) + dot := GetSimType[float32](DotProd, 32) + euclid := GetSimType[float32](Euclidean, 32) + + // Cosine / dot product: returned as-is (already higher-is-better). + require.InDelta(t, 0.9, cosine.similarityScore(0.9), 1e-6) + require.InDelta(t, -0.2, cosine.similarityScore(-0.2), 1e-6) + require.InDelta(t, 12.5, dot.similarityScore(12.5), 1e-6) + + // Euclidean: squared distance mapped to 1/(1+d), so a smaller distance yields a + // larger score (closer = better). + near := euclid.similarityScore(0.0) // distance 0 -> perfect match + mid := euclid.similarityScore(1.0) + far := euclid.similarityScore(9.0) + require.InDelta(t, 1.0, near, 1e-6) + require.InDelta(t, 0.5, mid, 1e-6) + require.InDelta(t, 0.1, far, 1e-6) + require.Greater(t, near, mid) + require.Greater(t, mid, far) +} diff --git a/tok/index/index.go b/tok/index/index.go index 1e981ef189e..050ec85e97e 100644 --- a/tok/index/index.go +++ b/tok/index/index.go @@ -147,6 +147,29 @@ type OptionalSearchOptions[T c.Float] interface { maxResults int, opts VectorIndexOptions[T]) ([]uint64, error) } +// ScoredSearchOptions extends search to also return a higher-is-better similarity +// score for each returned uid (positionally aligned). These power native hybrid +// search: the score is bound to a DQL value variable so vector results can be a +// fusion channel alongside BM25. Scores carry the same higher-is-better convention +// as BM25 (cosine/dot as-is; euclidean as 1/(1+dist)). +// +// SearchScored / SearchWithUidScored preserve the exact neighbor selection of +// Search / SearchWithUid (so scoring a plain query does not change its results), +// while the *Options* variants apply per-call ef/distance-threshold controls. +type ScoredSearchOptions[T c.Float] interface { + SearchScored(ctx context.Context, c CacheType, query []T, + maxResults int, filter SearchFilter[T]) ([]uint64, []float64, error) + + SearchWithUidScored(ctx context.Context, c CacheType, queryUid uint64, + maxResults int, filter SearchFilter[T]) ([]uint64, []float64, error) + + SearchWithOptionsScored(ctx context.Context, c CacheType, query []T, + maxResults int, opts VectorIndexOptions[T]) ([]uint64, []float64, error) + + SearchWithUidAndOptionsScored(ctx context.Context, c CacheType, queryUid uint64, + maxResults int, opts VectorIndexOptions[T]) ([]uint64, []float64, error) +} + // A Txn is an interface representation of a persistent storage transaction, // where multiple operations are performed on a database type Txn interface { diff --git a/tok/index/search_path.go b/tok/index/search_path.go index 7e24b7d068c..5387b71ee7f 100644 --- a/tok/index/search_path.go +++ b/tok/index/search_path.go @@ -12,6 +12,10 @@ type SearchPathResult struct { // The collection of nearest neighbors in sorted order after filtering // out neighbors that fail any Filter criteria. Neighbors []uint64 + // Distances holds the higher-is-better similarity score for each entry in + // Neighbors (positionally aligned). It is populated by scored searches and may + // be empty for callers that only need the neighbor uids. + Distances []float64 // The path from the start of search to the closest neighbor vector. Path []uint64 // A collection of captured named counters that occurred for the @@ -24,6 +28,7 @@ type SearchPathResult struct { func NewSearchPathResult() *SearchPathResult { return &SearchPathResult{ Neighbors: []uint64{}, + Distances: []float64{}, Path: []uint64{}, Metrics: make(map[string]uint64), } diff --git a/tok/tok.go b/tok/tok.go index c1da3e991d7..cb50b0a369e 100644 --- a/tok/tok.go +++ b/tok/tok.go @@ -50,6 +50,7 @@ const ( IdentBigFloat = 0xD IdentVFloat = 0xE IdentNGram = 0xF + IdentBM25 = 0x10 IdentCustom = 0x80 IdentDelimiter = 0x1f // ASCII 31 - Unit separator ) @@ -101,6 +102,7 @@ func init() { registerTokenizer(TermTokenizer{}) registerTokenizer(FullTextTokenizer{}) registerTokenizer(NGramTokenizer{}) + registerTokenizer(BM25Tokenizer{}) registerTokenizer(Sha256Tokenizer{}) setupBleve() } @@ -576,6 +578,47 @@ func (t FullTextTokenizer) Identifier() byte { return IdentFullText } func (t FullTextTokenizer) IsSortable() bool { return false } func (t FullTextTokenizer) IsLossy() bool { return true } +// BM25Tokenizer generates tokens for BM25 ranked text search. +// It uses the same pipeline as FullTextTokenizer (normalize, stopwords, stem) +// but preserves duplicates for term frequency counting. +type BM25Tokenizer struct{ lang string } + +func (t BM25Tokenizer) Name() string { return "bm25" } +func (t BM25Tokenizer) Type() string { return "string" } +func (t BM25Tokenizer) Tokens(v interface{}) ([]string, error) { + str, ok := v.(string) + if !ok || str == "" { + return []string{}, nil + } + lang := LangBase(t.lang) + tokens := fulltextAnalyzer.Analyze([]byte(str)) + tokens = filterStopwords(lang, tokens) + tokens = filterStemmers(lang, tokens) + // Return all tokens with duplicates preserved (for TF counting). + result := make([]string, 0, len(tokens)) + for _, t := range tokens { + result = append(result, string(t.Term)) + } + return result, nil +} +func (t BM25Tokenizer) Identifier() byte { return IdentBM25 } +func (t BM25Tokenizer) IsSortable() bool { return false } +func (t BM25Tokenizer) IsLossy() bool { return true } + +// TokensWithFrequency tokenizes the input and returns term frequencies and doc length. +func (t BM25Tokenizer) TokensWithFrequency(v interface{}, lang string) (map[string]uint32, uint32, error) { + tok := BM25Tokenizer{lang: lang} + allTokens, err := tok.Tokens(v) + if err != nil { + return nil, 0, err + } + termFreqs := make(map[string]uint32, len(allTokens)) + for _, t := range allTokens { + termFreqs[t]++ + } + return termFreqs, uint32(len(allTokens)), nil +} + // Sha256Tokenizer generates tokens for the sha256 hash part from string data. type Sha256Tokenizer struct{ _ string } diff --git a/tok/tok_test.go b/tok/tok_test.go index 4c95094e577..b9fbc4dd1a5 100644 --- a/tok/tok_test.go +++ b/tok/tok_test.go @@ -652,6 +652,146 @@ func TestNGramTokenizerNonStringInput(t *testing.T) { require.Equal(t, 0, len(tokens2), "Expected empty tokens for nil input") } +func TestBM25Tokenizer(t *testing.T) { + tokenizer, has := GetTokenizer("bm25") + require.True(t, has) + require.NotNil(t, tokenizer) + require.Equal(t, "bm25", tokenizer.Name()) + require.Equal(t, "string", tokenizer.Type()) + require.Equal(t, byte(IdentBM25), tokenizer.Identifier()) + require.True(t, tokenizer.IsLossy()) + require.False(t, tokenizer.IsSortable()) +} + +func TestBM25TokensPreservesDuplicates(t *testing.T) { + tok := BM25Tokenizer{lang: "en"} + tokens, err := tok.Tokens("fox fox fox dog") + require.NoError(t, err) + // "fox" should appear 3 times (duplicates preserved), "dog" once + foxCount := 0 + dogCount := 0 + for _, token := range tokens { + if token == "fox" { + foxCount++ + } + if token == "dog" { + dogCount++ + } + } + require.Equal(t, 3, foxCount, "Expected 3 occurrences of 'fox'") + require.Equal(t, 1, dogCount, "Expected 1 occurrence of 'dog'") +} + +func TestBM25TokensWithFrequency(t *testing.T) { + tok := BM25Tokenizer{} + termFreqs, docLen, err := tok.TokensWithFrequency("the quick brown fox fox fox", "en") + require.NoError(t, err) + // "the" is a stopword and should be removed + _, hasThe := termFreqs["the"] + require.False(t, hasThe, "'the' should be removed as stopword") + // "fox" should have tf=3 + require.Equal(t, uint32(3), termFreqs["fox"]) + // "quick" -> "quick" (stemmed) + require.Contains(t, termFreqs, "quick") + require.Equal(t, uint32(1), termFreqs["quick"]) + // "brown" -> "brown" (stemmed) + require.Contains(t, termFreqs, "brown") + require.Equal(t, uint32(1), termFreqs["brown"]) + // docLen should be total tokens after stopword removal + require.Equal(t, uint32(5), docLen) +} + +func TestBM25TokensEmpty(t *testing.T) { + tok := BM25Tokenizer{lang: "en"} + tokens, err := tok.Tokens("") + require.NoError(t, err) + require.Equal(t, 0, len(tokens)) + + termFreqs, docLen, err := tok.TokensWithFrequency("", "en") + require.NoError(t, err) + require.Equal(t, 0, len(termFreqs)) + require.Equal(t, uint32(0), docLen) +} + +func TestBM25TokensSingleWord(t *testing.T) { + tok := BM25Tokenizer{lang: "en"} + tokens, err := tok.Tokens("hello") + require.NoError(t, err) + require.Equal(t, 1, len(tokens)) + require.Equal(t, "hello", tokens[0]) +} + +func TestBM25TokensStemming(t *testing.T) { + tok := BM25Tokenizer{lang: "en"} + tokens, err := tok.Tokens("running jumping swimming") + require.NoError(t, err) + require.Equal(t, 3, len(tokens)) + require.Contains(t, tokens, "run") + require.Contains(t, tokens, "jump") + require.Contains(t, tokens, "swim") +} + +func TestGetBM25QueryTokens(t *testing.T) { + tokens, err := GetBM25QueryTokens([]string{"quick brown fox fox"}, "en") + require.NoError(t, err) + // Query tokens should be deduplicated + require.Equal(t, 3, len(tokens)) + // Each token should be encoded with the BM25 identifier prefix + for _, token := range tokens { + require.Equal(t, byte(IdentBM25), token[0], "Token should start with BM25 identifier") + } +} + +func TestGetBM25QueryTokensEmpty(t *testing.T) { + tokens, err := GetBM25QueryTokens([]string{""}, "en") + require.NoError(t, err) + require.Equal(t, 0, len(tokens)) +} + +func TestBM25TokenizerForLang(t *testing.T) { + tokenizer, has := GetTokenizer("bm25") + require.True(t, has) + langTok := GetTokenizerForLang(tokenizer, "de") + bm25Tok, ok := langTok.(BM25Tokenizer) + require.True(t, ok) + // German: "Katzen" -> "katz" (stemmed) + tokens, err := bm25Tok.Tokens("Katzen und Katzen") + require.NoError(t, err) + // "und" is a German stopword + katzCount := 0 + for _, token := range tokens { + if token == "katz" { + katzCount++ + } + } + require.Equal(t, 2, katzCount, "Expected 2 occurrences of stemmed 'katz'") +} + +func TestBM25AllStopwords(t *testing.T) { + tok := BM25Tokenizer{lang: "en"} + tokens, err := tok.Tokens("the a an is") + require.NoError(t, err) + require.Equal(t, 0, len(tokens)) + + termFreqs, docLen, err := tok.TokensWithFrequency("the a an is", "en") + require.NoError(t, err) + require.Equal(t, 0, len(termFreqs)) + require.Equal(t, uint32(0), docLen) +} + +func TestGetBM25QueryTokensAllStopwords(t *testing.T) { + tokens, err := GetBM25QueryTokens([]string{"the a an"}, "en") + require.NoError(t, err) + require.Equal(t, 0, len(tokens)) +} + +func TestGetBM25QueryTokensWrongArgCount(t *testing.T) { + _, err := GetBM25QueryTokens([]string{}, "en") + require.Error(t, err) + _, err = GetBM25QueryTokens([]string{"a", "b"}, "en") + require.Error(t, err) +} + func BenchmarkTermTokenizer(b *testing.B) { b.Skip() // tmp } diff --git a/tok/tokens.go b/tok/tokens.go index bda9a04e743..f089a3f4344 100644 --- a/tok/tokens.go +++ b/tok/tokens.go @@ -25,6 +25,8 @@ func GetTokenizerForLang(t Tokenizer, lang string) Tokenizer { // We must return a new instance because another goroutine might be calling this // with a different lang. return FullTextTokenizer{lang: lang} + case BM25Tokenizer: + return BM25Tokenizer{lang: lang} case TermTokenizer: return TermTokenizer{lang: lang} case ExactTokenizer: @@ -67,6 +69,29 @@ func GetNGramQueryTokens(funcArgs []string, lang string) ([]string, error) { return BuildNGramQueryTokens(funcArgs[0], NGramTokenizer{lang: lang}) } +// GetBM25QueryTokens tokenizes the query text using the fulltext pipeline, +// deduplicates, and encodes with the BM25 identifier prefix. +func GetBM25QueryTokens(funcArgs []string, lang string) ([]string, error) { + if l := len(funcArgs); l != 1 { + return nil, errors.Errorf("Function requires 1 arguments, but got %d", l) + } + tok := BM25Tokenizer{lang: lang} + allTokens, err := tok.Tokens(funcArgs[0]) + if err != nil { + return nil, err + } + // Deduplicate for query + seen := make(map[string]struct{}, len(allTokens)) + var unique []string + for _, t := range allTokens { + if _, ok := seen[t]; !ok { + seen[t] = struct{}{} + unique = append(unique, encodeToken(t, tok.Identifier())) + } + } + return unique, nil +} + // GetFullTextTokens returns the full-text tokens for the given value. func GetFullTextTokens(funcArgs []string, lang string) ([]string, error) { if l := len(funcArgs); l != 1 { diff --git a/worker/bm25wand.go b/worker/bm25wand.go new file mode 100644 index 00000000000..6b0881f3ddd --- /dev/null +++ b/worker/bm25wand.go @@ -0,0 +1,444 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package worker + +import ( + "container/heap" + "math" + "sort" + + "github.com/dgraph-io/dgraph/v25/posting" +) + +// wandBlockSize is the number of postings grouped into one logical block for +// Block-Max WAND upper bounds. The postings come from a standard Dgraph posting +// list (already resident in memory once loaded); these blocks exist only to give +// WAND per-block score bounds for pruning — they are not a storage format. +const wandBlockSize = 128 + +// termCursor is an in-memory cursor over one query term's posting list, +// materialized from the standard posting List as UID-ascending (uid, tf, docLen) +// entries. Document length travels with each posting, so scoring needs no separate +// lookup. Per-block max bounds drive Block-Max WAND pruning. +type termCursor struct { + postings []posting.BM25Posting + idf float64 + pos int + + // blockUBPre[i] is the pre-IDF BM25 upper bound for block i (max term + // frequency, min document length in the block). suffixUBPre[i] = max over + // j >= i of blockUBPre[j], for the remaining-list upper bound. + blockUBPre []float64 + suffixUBPre []float64 +} + +// ubPre computes the pre-IDF BM25 contribution upper bound for a block, using the +// block's maximum term frequency and minimum document length (the score is +// increasing in tf and decreasing in dl, so this is a safe upper bound). +func ubPre(maxTF, minDL uint32, k, b, avgDL float64) float64 { + if avgDL <= 0 { + avgDL = 1 + } + tf := float64(maxTF) + dl := float64(minDL) + denom := k*(1-b+b*dl/avgDL) + tf + if denom <= 0 { + return 0 + } + return (k + 1) * tf / denom +} + +// newTermCursor builds a cursor and precomputes its per-block upper bounds. +func newTermCursor(postings []posting.BM25Posting, idf, k, b, avgDL float64) *termCursor { + c := &termCursor{postings: postings, idf: idf} + numBlocks := (len(postings) + wandBlockSize - 1) / wandBlockSize + c.blockUBPre = make([]float64, numBlocks) + for blk := 0; blk < numBlocks; blk++ { + start := blk * wandBlockSize + end := start + wandBlockSize + if end > len(postings) { + end = len(postings) + } + var maxTF uint32 + minDL := uint32(math.MaxUint32) + for i := start; i < end; i++ { + if postings[i].TF > maxTF { + maxTF = postings[i].TF + } + dl := postings[i].DocLen + if dl == 0 { + dl = 1 + } + if dl < minDL { + minDL = dl + } + } + c.blockUBPre[blk] = ubPre(maxTF, minDL, k, b, avgDL) + } + c.suffixUBPre = make([]float64, numBlocks) + var running float64 + for blk := numBlocks - 1; blk >= 0; blk-- { + if c.blockUBPre[blk] > running { + running = c.blockUBPre[blk] + } + c.suffixUBPre[blk] = running + } + return c +} + +func (c *termCursor) exhausted() bool { return c.pos >= len(c.postings) } + +func (c *termCursor) currentDoc() uint64 { + if c.exhausted() { + return math.MaxUint64 + } + return c.postings[c.pos].Uid +} + +func (c *termCursor) currentTF() uint32 { + if c.exhausted() { + return 0 + } + return c.postings[c.pos].TF +} + +func (c *termCursor) currentDocLen() uint32 { + if c.exhausted() { + return 0 + } + return c.postings[c.pos].DocLen +} + +// remainingUB returns the IDF-weighted upper-bound score over the remainder of the +// list from the current position. +func (c *termCursor) remainingUB() float64 { + if c.exhausted() || len(c.suffixUBPre) == 0 { + return 0 + } + blk := c.pos / wandBlockSize + if blk >= len(c.suffixUBPre) { + return 0 + } + return c.idf * c.suffixUBPre[blk] +} + +// next advances by one posting. +func (c *termCursor) next() bool { + c.pos++ + return !c.exhausted() +} + +// skipTo advances to the first posting with UID >= target. +func (c *termCursor) skipTo(target uint64) bool { + if c.exhausted() { + return false + } + if c.postings[c.pos].Uid >= target { + return true + } + rel := sort.Search(len(c.postings)-c.pos, func(i int) bool { + return c.postings[c.pos+i].Uid >= target + }) + c.pos += rel + return !c.exhausted() +} + +// skipToWithBMW is skipTo with Block-Max WAND pruning: blocks whose upper bound +// combined with otherUB cannot beat theta are skipped wholesale. +func (c *termCursor) skipToWithBMW(target uint64, theta, otherUB float64) bool { + if !c.skipTo(target) { + return false + } + for !c.exhausted() { + blk := c.pos / wandBlockSize + if c.idf*c.blockUBPre[blk]+otherUB > theta { + return true + } + // This block can't produce a winner; jump to the start of the next block. + c.pos = (blk + 1) * wandBlockSize + } + return false +} + +// scoredDoc holds a UID and its BM25 score for the min-heap. +type scoredDoc struct { + uid uint64 + score float64 +} + +// topKHeap is a min-heap of scored documents for top-k tracking. +type topKHeap struct { + docs []scoredDoc + k int +} + +func (h *topKHeap) Len() int { return len(h.docs) } +func (h *topKHeap) Less(i, j int) bool { + if h.docs[i].score != h.docs[j].score { + return h.docs[i].score < h.docs[j].score + } + // Among equal scores, treat the higher UID as "smaller" so it sits at the root and + // is the eviction victim first. This keeps the lowest UID on score ties, matching + // the (score desc, UID asc) order that sorted() and scoreAllDocs produce, so a + // first:k WAND query returns the same set as the exhaustive path. Primary ordering + // stays score-ascending, so threshold() still reports the true minimum score and + // WAND pivot pruning is unaffected. + return h.docs[i].uid > h.docs[j].uid +} +func (h *topKHeap) Swap(i, j int) { h.docs[i], h.docs[j] = h.docs[j], h.docs[i] } +func (h *topKHeap) Push(x interface{}) { h.docs = append(h.docs, x.(scoredDoc)) } +func (h *topKHeap) Pop() interface{} { + old := h.docs + n := len(old) + item := old[n-1] + h.docs = old[:n-1] + return item +} + +// threshold returns the minimum score in the heap (the score to beat). +func (h *topKHeap) threshold() float64 { + if len(h.docs) < h.k { + return 0 + } + return h.docs[0].score +} + +// tryPush adds a doc if it beats the current threshold. +func (h *topKHeap) tryPush(uid uint64, score float64) { + if len(h.docs) < h.k { + heap.Push(h, scoredDoc{uid: uid, score: score}) + return + } + if score > h.docs[0].score { + h.docs[0] = scoredDoc{uid: uid, score: score} + heap.Fix(h, 0) + } +} + +// sorted returns all docs sorted by score descending, then UID ascending. +func (h *topKHeap) sorted() []scoredDoc { + result := make([]scoredDoc, len(h.docs)) + copy(result, h.docs) + sort.Slice(result, func(i, j int) bool { + if result[i].score != result[j].score { + return result[i].score > result[j].score + } + return result[i].uid < result[j].uid + }) + return result +} + +// bm25TopK returns how many top-scored documents the WAND search should retain for a +// first/offset query window. With a first limit it is first+offset (so the offset can +// be dropped afterward while still bounding work and memory to the window); with no +// first limit it is 0, meaning every matching document is scored. Returning first+offset +// rather than 0 whenever an offset is present is what keeps a deep-paginated query from +// materializing and scoring the entire corpus. +func bm25TopK(first, offset int) int { + if first <= 0 { + return 0 + } + return first + offset +} + +// bm25PaginateScored slices score-descending results to the [offset, offset+first) +// window. first <= 0 means no upper bound. It clamps offset to the slice length so an +// offset past the end yields an empty result instead of panicking. +func bm25PaginateScored(results []scoredDoc, first, offset int) []scoredDoc { + if offset > len(results) { + offset = len(results) + } + results = results[offset:] + if first > 0 && first < len(results) { + results = results[:first] + } + return results +} + +// bm25Score computes the BM25 contribution of a single term occurrence. +func bm25Score(idf, tf, dl, avgDL, k, b float64) float64 { + if avgDL <= 0 { + avgDL = 1 + } + if dl <= 0 { + dl = 1 + } + return idf * (k + 1) * tf / (k*(1-b+b*dl/avgDL) + tf) +} + +// wandSearch performs a WAND / Block-Max WAND top-k BM25 search over standard +// posting lists. queryTokens must already carry the BM25 tokenizer identifier +// byte. getList reads a posting list for a key. If topK <= 0, every matching +// document is scored (no early termination). +func wandSearch(getList func(key []byte) (*posting.List, error), attr string, readTs uint64, + queryTokens []string, k, b, avgDL, N float64, topK int, + filterSet map[uint64]struct{}, useBMW bool) ([]scoredDoc, error) { + + var cursors []*termCursor + for _, token := range queryTokens { + postings, err := posting.ReadBM25TermPostings(getList, attr, token, readTs) + if err != nil { + return nil, err + } + df := uint64(len(postings)) + if df == 0 { + continue + } + // N comes from bucketed stats and df from the term's posting list; if stats + // ever lag the postings, clamp N >= df for this term so the smoothed IDF + // stays non-negative and finite instead of producing a negative/NaN score. + dfN := float64(df) + nDocs := N + if nDocs < dfN { + nDocs = dfN + } + idf := math.Log1p((nDocs - dfN + 0.5) / (dfN + 0.5)) + cursors = append(cursors, newTermCursor(postings, idf, k, b, avgDL)) + } + + if len(cursors) == 0 { + return nil, nil + } + + if topK <= 0 { + return scoreAllDocs(cursors, k, b, avgDL, filterSet), nil + } + return wandTopK(cursors, k, b, avgDL, topK, filterSet, useBMW), nil +} + +// wandTopK runs the WAND / Block-Max WAND main loop over prepared cursors and +// returns the top-k documents sorted by score descending. It is the core scoring +// loop, separated from posting-list I/O so it can be exercised directly. +func wandTopK(cursors []*termCursor, k, b, avgDL float64, topK int, + filterSet map[uint64]struct{}, useBMW bool) []scoredDoc { + + h := &topKHeap{k: topK} + heap.Init(h) + + for { + // Drop exhausted cursors. + active := cursors[:0] + for _, c := range cursors { + if !c.exhausted() { + active = append(active, c) + } + } + cursors = active + if len(cursors) == 0 { + break + } + + // Sort cursors by current document ascending. + sort.Slice(cursors, func(i, j int) bool { + return cursors[i].currentDoc() < cursors[j].currentDoc() + }) + + theta := h.threshold() + + // Find pivot: accumulate upper bounds until they exceed theta. + var sumUB float64 + pivot := -1 + var pivotDoc uint64 + for i, c := range cursors { + sumUB += c.remainingUB() + if sumUB > theta && pivot == -1 { + pivot = i + pivotDoc = c.currentDoc() + } + } + if pivot == -1 { + break // sum of all upper bounds can't beat theta + } + + // Advance all cursors before the pivot up to pivotDoc. + allAtPivot := true + for i := 0; i < pivot; i++ { + if cursors[i].currentDoc() < pivotDoc { + var ok bool + if useBMW { + otherUB := sumUB - cursors[i].remainingUB() + ok = cursors[i].skipToWithBMW(pivotDoc, theta, otherUB) + } else { + ok = cursors[i].skipTo(pivotDoc) + } + if !ok { + allAtPivot = false + break + } + if cursors[i].currentDoc() != pivotDoc { + allAtPivot = false + } + } + } + if !allAtPivot { + continue + } + + // Score the pivot document. + if filterSet != nil { + if _, ok := filterSet[pivotDoc]; !ok { + for _, c := range cursors { + if c.currentDoc() == pivotDoc { + c.next() + } + } + continue + } + } + + var score float64 + for _, c := range cursors { + if c.currentDoc() == pivotDoc { + dl := float64(c.currentDocLen()) + score += bm25Score(c.idf, float64(c.currentTF()), dl, avgDL, k, b) + } + } + h.tryPush(pivotDoc, score) + + for _, c := range cursors { + if c.currentDoc() == pivotDoc { + c.next() + } + } + } + + return h.sorted() +} + +// scoreAllDocs scores every matching document without early termination. Used when +// no top-k limit is specified. +func scoreAllDocs(cursors []*termCursor, k, b, avgDL float64, + filterSet map[uint64]struct{}) []scoredDoc { + + scores := make(map[uint64]float64) + + for _, c := range cursors { + for !c.exhausted() { + uid := c.currentDoc() + if filterSet != nil { + if _, ok := filterSet[uid]; !ok { + c.next() + continue + } + } + scores[uid] += bm25Score(c.idf, float64(c.currentTF()), float64(c.currentDocLen()), + avgDL, k, b) + c.next() + } + } + + results := make([]scoredDoc, 0, len(scores)) + for uid, s := range scores { + results = append(results, scoredDoc{uid: uid, score: s}) + } + sort.Slice(results, func(i, j int) bool { + if results[i].score != results[j].score { + return results[i].score > results[j].score + } + return results[i].uid < results[j].uid + }) + return results +} diff --git a/worker/bm25wand_bench_test.go b/worker/bm25wand_bench_test.go new file mode 100644 index 00000000000..38bf09d707d --- /dev/null +++ b/worker/bm25wand_bench_test.go @@ -0,0 +1,352 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package worker + +// Q31 WAND effectiveness benchmarks. Synthetic corpora (10k and 100k docs) with +// Zipfian term selectivity — a rare term (df ~0.1%), a mid term (df ~5%), and a +// stopword-like term (df ~60%) — are fed directly into wandTopK / scoreAllDocs, +// mirroring what wandSearch does after ReadBM25TermPostings materializes each +// term's posting list. The matrix covers 1–3 term queries, topK in {10, 100} for +// WAND (BMW on and off) against the topK=0 exhaustive scoreAllDocs path, plus the +// deep-offset shape (first:10 offset:10000 => topK=10010 via bm25TopK) so the +// early-termination crossover is visible. +// +// Each benchmark reports two domain metrics alongside ns/op: +// - docs_scored/op: documents fully scored (tryPush'd for WAND, accumulated for +// scoreAllDocs), counted by a replica of the wandTopK loop since the product +// code exposes no counter. The replica is checked bit-identical against the +// real wandTopK before timing, so it cannot drift silently. The count is +// deterministic per config, computed off the clock, and reported after the +// timed loop (b.ResetTimer clears metrics reported before it). +// - postings/op: total materialized postings across the query's cursors — the +// work floor that full materialization imposes regardless of pruning. +// +// Cursor construction (newTermCursor block-bound precompute) is inside the timed +// loop, matching wandSearch, which rebuilds cursors per query. + +import ( + "container/heap" + "math" + "math/rand" + "sort" + "strconv" + "testing" + + "github.com/dgraph-io/dgraph/v25/posting" +) + +const benchSeed = 20260718 + +// benchTerm is one query term's materialized posting list plus its smoothed IDF, +// exactly as wandSearch computes it. +type benchTerm struct { + postings []posting.BM25Posting + idf float64 +} + +// benchCorpus is a deterministic synthetic corpus: per-doc lengths shared across +// terms, and three terms tiered by document frequency. +type benchCorpus struct { + numDocs int + avgDL float64 + terms map[string]benchTerm +} + +var benchCorpora = map[int]*benchCorpus{} + +// getBenchCorpus builds (once per size) a corpus of numDocs documents with +// uniform doc lengths in [5, 80] and Zipf-distributed term frequencies, then +// materializes posting lists for the three selectivity tiers: +// +// rare ~0.1% df, mid ~5% df, stop ~60% df. +// +// Membership is an independent coin flip per (doc, term) at the tier's rate, so +// postings interleave uniformly across the UID space (UID-ascending, as real +// posting lists are). Benchmarks share corpora; cursors are rebuilt per iteration +// because they are stateful. +func getBenchCorpus(numDocs int) *benchCorpus { + if c, ok := benchCorpora[numDocs]; ok { + return c + } + rng := rand.New(rand.NewSource(benchSeed + int64(numDocs))) + zipfTF := rand.NewZipf(rng, 1.5, 1, 9) // tf-1 in 0..9, Zipf-skewed toward 0 + + docLens := make([]uint32, numDocs+1) // 1-based UIDs + var totalLen float64 + for uid := 1; uid <= numDocs; uid++ { + docLens[uid] = uint32(5 + rng.Intn(76)) + totalLen += float64(docLens[uid]) + } + avgDL := totalLen / float64(numDocs) + + c := &benchCorpus{numDocs: numDocs, avgDL: avgDL, terms: map[string]benchTerm{}} + for _, tier := range []struct { + name string + rate float64 + }{ + {"rare", 0.001}, + {"mid", 0.05}, + {"stop", 0.60}, + } { + var ps []posting.BM25Posting + for uid := 1; uid <= numDocs; uid++ { + if rng.Float64() >= tier.rate { + continue + } + ps = append(ps, posting.BM25Posting{ + Uid: uint64(uid), + TF: uint32(1 + zipfTF.Uint64()), + DocLen: docLens[uid], + }) + } + // Smoothed IDF exactly as wandSearch computes it (N >= df by construction). + df := float64(len(ps)) + idf := math.Log1p((float64(numDocs) - df + 0.5) / (df + 0.5)) + c.terms[tier.name] = benchTerm{postings: ps, idf: idf} + } + benchCorpora[numDocs] = c + return c +} + +// buildBenchCursors materializes fresh cursors for the named terms, as wandSearch +// does per query. +func buildBenchCursors(c *benchCorpus, termNames []string, k, b float64) []*termCursor { + cs := make([]*termCursor, 0, len(termNames)) + for _, name := range termNames { + t := c.terms[name] + if len(t.postings) == 0 { + continue + } + cs = append(cs, newTermCursor(t.postings, t.idf, k, b, c.avgDL)) + } + return cs +} + +func benchTotalPostings(c *benchCorpus, termNames []string) int { + var n int + for _, name := range termNames { + n += len(c.terms[name].postings) + } + return n +} + +// wandTopKCounting is a line-for-line replica of wandTopK (filterSet elided — +// the benchmarks pass nil) that additionally counts pivot documents fully scored. +// The product code exposes no counter, so the count lives in this harness copy; +// benchmarks verify its output bit-identical to the real wandTopK before timing. +func wandTopKCounting(cursors []*termCursor, k, b, avgDL float64, topK int, + useBMW bool) (results []scoredDoc, docsScored int) { + + h := &topKHeap{k: topK} + heap.Init(h) + + for { + active := cursors[:0] + for _, c := range cursors { + if !c.exhausted() { + active = append(active, c) + } + } + cursors = active + if len(cursors) == 0 { + break + } + + sort.Slice(cursors, func(i, j int) bool { + return cursors[i].currentDoc() < cursors[j].currentDoc() + }) + + theta := h.threshold() + + var sumUB float64 + pivot := -1 + var pivotDoc uint64 + for i, c := range cursors { + sumUB += c.remainingUB() + if sumUB > theta && pivot == -1 { + pivot = i + pivotDoc = c.currentDoc() + } + } + if pivot == -1 { + break + } + + allAtPivot := true + for i := 0; i < pivot; i++ { + if cursors[i].currentDoc() < pivotDoc { + var ok bool + if useBMW { + otherUB := sumUB - cursors[i].remainingUB() + ok = cursors[i].skipToWithBMW(pivotDoc, theta, otherUB) + } else { + ok = cursors[i].skipTo(pivotDoc) + } + if !ok { + allAtPivot = false + break + } + if cursors[i].currentDoc() != pivotDoc { + allAtPivot = false + } + } + } + if !allAtPivot { + continue + } + + var score float64 + for _, c := range cursors { + if c.currentDoc() == pivotDoc { + dl := float64(c.currentDocLen()) + score += bm25Score(c.idf, float64(c.currentTF()), dl, avgDL, k, b) + } + } + docsScored++ + h.tryPush(pivotDoc, score) + + for _, c := range cursors { + if c.currentDoc() == pivotDoc { + c.next() + } + } + } + + return h.sorted(), docsScored +} + +// benchWandDocsScored computes docs_scored for one config via the counting +// replica and guards the replica against drift from the real wandTopK +// (bit-identical uids and scores). Called before the timed loop; the caller +// must ReportMetric AFTER the loop, because b.ResetTimer clears extra metrics. +func benchWandDocsScored(b *testing.B, c *benchCorpus, termNames []string, + k, bp float64, topK int, useBMW bool) int { + + want := wandTopK(buildBenchCursors(c, termNames, k, bp), k, bp, c.avgDL, + topK, nil, useBMW) + got, docsScored := wandTopKCounting(buildBenchCursors(c, termNames, k, bp), + k, bp, c.avgDL, topK, useBMW) + if len(got) != len(want) { + b.Fatalf("counting replica drifted from wandTopK: len %d != %d", len(got), len(want)) + } + for i := range want { + if got[i].uid != want[i].uid || + math.Float64bits(got[i].score) != math.Float64bits(want[i].score) { + b.Fatalf("counting replica drifted from wandTopK at rank %d: "+ + "got (%d, %v) want (%d, %v)", i, got[i].uid, got[i].score, + want[i].uid, want[i].score) + } + } + return docsScored +} + +// BenchmarkWandEffectiveness is the Q31 matrix: corpus size x query shape x +// {exhaustive, wand/bmw x topK}. scoreAll is the topK=0 path a query with no +// first: limit takes; wand/bmw topK=10/100 are the first:10/first:100 windows. +func BenchmarkWandEffectiveness(b *testing.B) { + const k, bp = 1.2, 0.75 + + queries := []struct { + name string + terms []string + }{ + {"q=1term-rare", []string{"rare"}}, + {"q=1term-stop", []string{"stop"}}, + {"q=2term-rare+stop", []string{"rare", "stop"}}, + {"q=3term-rare+mid+stop", []string{"rare", "mid", "stop"}}, + } + + for _, size := range []struct { + name string + docs int + }{ + {"docs=10k", 10_000}, + {"docs=100k", 100_000}, + } { + c := getBenchCorpus(size.docs) + for _, q := range queries { + prefix := size.name + "/" + q.name + + // topK=0: the exhaustive scoreAllDocs path (no first: limit). + b.Run(prefix+"/algo=scoreAll/topK=0", func(b *testing.B) { + full := scoreAllDocs(buildBenchCursors(c, q.terms, k, bp), k, bp, + c.avgDL, nil) + b.ResetTimer() + for i := 0; i < b.N; i++ { + cursors := buildBenchCursors(c, q.terms, k, bp) + _ = scoreAllDocs(cursors, k, bp, c.avgDL, nil) + } + // After the loop: ResetTimer clears metrics reported before it. + b.ReportMetric(float64(len(full)), "docs_scored/op") + b.ReportMetric(float64(benchTotalPostings(c, q.terms)), "postings/op") + }) + + for _, algo := range []struct { + name string + useBMW bool + }{ + {"algo=wand", false}, + {"algo=bmw", true}, + } { + for _, topK := range []int{10, 100} { + name := prefix + "/" + algo.name + "/topK=" + strconv.Itoa(topK) + useBMW := algo.useBMW + topK := topK + b.Run(name, func(b *testing.B) { + docsScored := benchWandDocsScored(b, c, q.terms, k, bp, topK, useBMW) + b.ResetTimer() + for i := 0; i < b.N; i++ { + cursors := buildBenchCursors(c, q.terms, k, bp) + _ = wandTopK(cursors, k, bp, c.avgDL, topK, nil, useBMW) + } + b.ReportMetric(float64(docsScored), "docs_scored/op") + b.ReportMetric(float64(benchTotalPostings(c, q.terms)), "postings/op") + }) + } + } + } + } +} + +// BenchmarkWandDeepOffset is the deep-pagination shape on the 100k corpus: +// first:10 offset:0 retains topK=10, while first:10 offset:10000 retains +// topK=10010 (bm25TopK), so the heap threshold stays 0 for the first 10010 docs +// and early termination degrades toward exhaustive. +func BenchmarkWandDeepOffset(b *testing.B) { + const k, bp = 1.2, 0.75 + c := getBenchCorpus(100_000) + terms := []string{"rare", "mid", "stop"} + + for _, cfg := range []struct { + name string + first int + off int + }{ + {"first=10/offset=0", 10, 0}, + {"first=10/offset=10000", 10, 10000}, + } { + topK := bm25TopK(cfg.first, cfg.off) + for _, algo := range []struct { + name string + useBMW bool + }{ + {"algo=wand", false}, + {"algo=bmw", true}, + } { + useBMW := algo.useBMW + b.Run(cfg.name+"/"+algo.name+"/topK="+strconv.Itoa(topK), func(b *testing.B) { + docsScored := benchWandDocsScored(b, c, terms, k, bp, topK, useBMW) + b.ResetTimer() + for i := 0; i < b.N; i++ { + cursors := buildBenchCursors(c, terms, k, bp) + _ = wandTopK(cursors, k, bp, c.avgDL, topK, nil, useBMW) + } + b.ReportMetric(float64(docsScored), "docs_scored/op") + b.ReportMetric(float64(benchTotalPostings(c, terms)), "postings/op") + }) + } + } +} diff --git a/worker/bm25wand_test.go b/worker/bm25wand_test.go new file mode 100644 index 00000000000..2007cb55fc8 --- /dev/null +++ b/worker/bm25wand_test.go @@ -0,0 +1,593 @@ +/* + * SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package worker + +import ( + "container/heap" + "math" + "math/rand" + "sort" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/posting" +) + +func TestTopKHeapBasic(t *testing.T) { + h := &topKHeap{k: 3} + heap.Init(h) + + require.Equal(t, 0.0, h.threshold()) + + h.tryPush(1, 5.0) + h.tryPush(2, 3.0) + require.Equal(t, 0.0, h.threshold()) // not full yet + + h.tryPush(3, 7.0) + require.InEpsilon(t, 3.0, h.threshold(), 1e-9) // full, min is 3.0 + + h.tryPush(4, 4.0) + require.InEpsilon(t, 4.0, h.threshold(), 1e-9) // 3.0 evicted, min is now 4.0 + + // 2.0 shouldn't be accepted. + h.tryPush(5, 2.0) + require.InEpsilon(t, 4.0, h.threshold(), 1e-9) + + sorted := h.sorted() + require.Len(t, sorted, 3) + require.Equal(t, uint64(3), sorted[0].uid) // highest score (7.0) + require.Equal(t, uint64(1), sorted[1].uid) // 5.0 + require.Equal(t, uint64(4), sorted[2].uid) // 4.0 +} + +func TestTopKHeapTieBreaking(t *testing.T) { + h := &topKHeap{k: 5} + heap.Init(h) + + // Same score, different UIDs — should sort by UID ascending. + h.tryPush(10, 5.0) + h.tryPush(5, 5.0) + h.tryPush(15, 5.0) + + sorted := h.sorted() + require.Equal(t, uint64(5), sorted[0].uid) + require.Equal(t, uint64(10), sorted[1].uid) + require.Equal(t, uint64(15), sorted[2].uid) +} + +// TestTopKHeapTieBreakEviction pins the boundary case the randomized brute-force +// tests deliberately skip (they treat tied-boundary uids as interchangeable). Docs +// arrive UID-ascending, as the WAND pivot advances. With k=2, two docs tied at score +// 1.0 (uid 3, uid 4) then a strictly higher doc (uid 5), the correct top-2 by +// (score desc, UID asc) is {5, 3}: among the tied pair uid3 (lower) outranks uid4, +// matching the no-limit scoreAllDocs path. A score-only heap evicts the lowest-UID +// root and wrongly returns {5, 4}. +func TestTopKHeapTieBreakEviction(t *testing.T) { + h := &topKHeap{k: 2} + heap.Init(h) + h.tryPush(3, 1.0) + h.tryPush(4, 1.0) + h.tryPush(5, 2.0) + + sorted := h.sorted() + uids := make([]uint64, len(sorted)) + for i, d := range sorted { + uids[i] = d.uid + } + require.Equal(t, []uint64{5, 3}, uids, + "WAND top-k must keep the lowest UID among score ties, matching scoreAllDocs") + + // The eviction threshold must remain the true minimum score after the tie-break. + require.InEpsilon(t, 1.0, h.threshold(), 1e-9) +} + +// TestWandTieBreakEndToEndMatchesFullScan exercises the full retrieval path (not just +// the heap) on a score tie that triggers eviction. Single term: uid10 and uid20 tie at +// a lower score (same tf/docLen), uid30 scores higher (larger tf). Processed +// UID-ascending, the two ties fill a k=2 heap, then uid30 evicts one — and it must +// evict the higher-UID tie (uid20), so first:2 returns {30, 10}, exactly the +// exhaustive scoreAllDocs top-2. This is the boundary TestWandMatchesBruteForce skips. +func TestWandTieBreakEndToEndMatchesFullScan(t *testing.T) { + uidsOf := func(ds []scoredDoc) []uint64 { + out := make([]uint64, len(ds)) + for i, d := range ds { + out[i] = d.uid + } + return out + } + ps := []posting.BM25Posting{ + {Uid: 10, TF: 2, DocLen: 5}, + {Uid: 20, TF: 2, DocLen: 5}, // ties with uid10 + {Uid: 30, TF: 8, DocLen: 5}, // strictly higher score + } + k, b, avgDL, idf := 1.2, 0.75, 5.0, 1.5 + build := func() []*termCursor { return []*termCursor{newTermCursor(ps, idf, k, b, avgDL)} } + + full := scoreAllDocs(build(), k, b, avgDL, nil) + require.Equal(t, []uint64{30, 10, 20}, uidsOf(full), "full scan: score desc, uid asc") + + for _, useBMW := range []bool{false, true} { + got := wandTopK(build(), k, b, avgDL, 2, nil, useBMW) + require.Equalf(t, []uint64{30, 10}, uidsOf(got), + "bmw=%v: first:k WAND must keep the same docs as full-scan top-k on ties", useBMW) + } +} + +func TestBm25TopK(t *testing.T) { + // No first limit: score every matching document (0 means "no early termination"). + require.Equal(t, 0, bm25TopK(0, 0)) + require.Equal(t, 0, bm25TopK(0, 100)) + + // With a first limit, WAND must retain first+offset documents so the offset can be + // dropped afterward — NOT 0 (which would fall back to scoring the entire corpus + // just because an offset was supplied, the memory blow-up this guards against). + require.Equal(t, 10, bm25TopK(10, 0)) + require.Equal(t, 15, bm25TopK(10, 5)) + require.Equal(t, 1001, bm25TopK(1, 1000)) +} + +func TestBm25PaginateScored(t *testing.T) { + mk := func(uids ...uint64) []scoredDoc { + out := make([]scoredDoc, len(uids)) + for i, u := range uids { + out[i] = scoredDoc{uid: u, score: float64(len(uids) - i)} // already score-descending + } + return out + } + ids := func(ds []scoredDoc) []uint64 { + out := make([]uint64, len(ds)) + for i, d := range ds { + out[i] = d.uid + } + return out + } + + full := mk(1, 2, 3, 4, 5) + require.Equal(t, []uint64{1, 2, 3, 4, 5}, ids(bm25PaginateScored(full, 0, 0))) + require.Equal(t, []uint64{1, 2}, ids(bm25PaginateScored(mk(1, 2, 3, 4, 5), 2, 0))) + require.Equal(t, []uint64{3, 4}, ids(bm25PaginateScored(mk(1, 2, 3, 4, 5), 2, 2))) + require.Equal(t, []uint64{4, 5}, ids(bm25PaginateScored(mk(1, 2, 3, 4, 5), 10, 3))) + // Offset past the end yields nothing rather than panicking. + require.Empty(t, bm25PaginateScored(mk(1, 2, 3), 2, 10)) +} + +func TestBm25ScoreFunction(t *testing.T) { + k, b := 1.2, 0.75 + avgDL := 10.0 + + // idf * (k+1) * tf / (k*(1-b+b*dl/avgDL) + tf) + idf := 1.5 + tf := 3.0 + dl := 10.0 + + expected := idf * (k + 1) * tf / (k*(1-b+b*dl/avgDL) + tf) + got := bm25Score(idf, tf, dl, avgDL, k, b) + require.InEpsilon(t, expected, got, 1e-9) + + // With b=0: no length normalization. + expected0 := idf * (k + 1) * tf / (k + tf) + got0 := bm25Score(idf, tf, dl, avgDL, k, 0) + require.InEpsilon(t, expected0, got0, 1e-9) + + // Score should be positive for positive inputs. + require.Greater(t, bm25Score(1.0, 1.0, 5.0, 10.0, k, b), 0.0) + + // Higher tf should produce higher score (same dl). + s1 := bm25Score(idf, 1.0, dl, avgDL, k, b) + s3 := bm25Score(idf, 3.0, dl, avgDL, k, b) + require.Greater(t, s3, s1) + + // Shorter doc should score higher (same tf). + sShort := bm25Score(idf, tf, 5.0, avgDL, k, b) + sLong := bm25Score(idf, tf, 20.0, avgDL, k, b) + require.Greater(t, sShort, sLong) +} + +func TestBm25ScoreNaN(t *testing.T) { + // Ensure no NaN/Inf for edge-case inputs. + score := bm25Score(0.5, 1.0, 0.0, 10.0, 1.2, 0.75) + require.False(t, math.IsNaN(score)) + require.False(t, math.IsInf(score, 0)) + require.Greater(t, score, 0.0) +} + +// brute force scores every doc across all cursors (ground truth for WAND). When +// filterSet is non-nil, only documents in it are scored — mirroring @filter(bm25(...)). +func bruteForceTopK(termPostings [][]posting.BM25Posting, idfs []float64, + k, b, avgDL float64, topK int) []scoredDoc { + return bruteForceTopKFiltered(termPostings, idfs, k, b, avgDL, topK, nil) +} + +func bruteForceTopKFiltered(termPostings [][]posting.BM25Posting, idfs []float64, + k, b, avgDL float64, topK int, filterSet map[uint64]struct{}) []scoredDoc { + scores := map[uint64]float64{} + for ti, ps := range termPostings { + for _, p := range ps { + if filterSet != nil { + if _, ok := filterSet[p.Uid]; !ok { + continue + } + } + scores[p.Uid] += bm25Score(idfs[ti], float64(p.TF), float64(p.DocLen), avgDL, k, b) + } + } + out := make([]scoredDoc, 0, len(scores)) + for uid, s := range scores { + out = append(out, scoredDoc{uid: uid, score: s}) + } + sort.Slice(out, func(i, j int) bool { + if out[i].score != out[j].score { + return out[i].score > out[j].score + } + return out[i].uid < out[j].uid + }) + if topK > 0 && len(out) > topK { + out = out[:topK] + } + return out +} + +// TestWandFilteredMatchesBruteForce checks that WAND/Block-Max WAND and the +// score-all path honor a filter set identically to exhaustive filtered scoring. The +// filter must never change which documents or scores are produced (only which are +// considered), so WAND pruning driven by a threshold built from filtered-in documents +// must still be sound. +func TestWandFilteredMatchesBruteForce(t *testing.T) { + rng := rand.New(rand.NewSource(7)) + k, b, avgDL := 1.2, 0.75, 9.0 + + for trial := 0; trial < 200; trial++ { + numTerms := 1 + rng.Intn(4) + termPostings := make([][]posting.BM25Posting, numTerms) + idfs := make([]float64, numTerms) + allUids := map[uint64]bool{} + for ti := 0; ti < numTerms; ti++ { + n := rng.Intn(400) + seen := map[uint64]bool{} + var ps []posting.BM25Posting + for j := 0; j < n; j++ { + uid := uint64(1 + rng.Intn(500)) + if seen[uid] { + continue + } + seen[uid] = true + ps = append(ps, posting.BM25Posting{ + Uid: uid, + TF: uint32(1 + rng.Intn(10)), + DocLen: uint32(1 + rng.Intn(30)), + }) + allUids[uid] = true + } + sort.Slice(ps, func(i, j int) bool { return ps[i].Uid < ps[j].Uid }) + termPostings[ti] = ps + idfs[ti] = 0.5 + rng.Float64()*2 + } + + // Random filter subset (may be empty). + filterSet := map[uint64]struct{}{} + for uid := range allUids { + if rng.Intn(2) == 0 { + filterSet[uid] = struct{}{} + } + } + + build := func() []*termCursor { + cs := make([]*termCursor, 0, numTerms) + for ti, ps := range termPostings { + if len(ps) == 0 { + continue + } + cs = append(cs, newTermCursor(ps, idfs[ti], k, b, avgDL)) + } + return cs + } + + // score-all path with filter must reproduce the full filtered ranking exactly. + wantAll := bruteForceTopKFiltered(termPostings, idfs, k, b, avgDL, 0, filterSet) + gotAll := scoreAllDocs(build(), k, b, avgDL, filterSet) + require.Lenf(t, gotAll, len(wantAll), "trial %d filtered score-all len", trial) + for i := range wantAll { + require.InEpsilonf(t, wantAll[i].score, gotAll[i].score, 1e-9, + "trial %d filtered score-all rank %d score", trial, i) + } + + // top-k WAND/BMW with filter must match the filtered top-k scores. + topK := 1 + rng.Intn(8) + want := bruteForceTopKFiltered(termPostings, idfs, k, b, avgDL, topK, filterSet) + wantPlus := bruteForceTopKFiltered(termPostings, idfs, k, b, avgDL, topK+1, filterSet) + for _, useBMW := range []bool{false, true} { + got := wandTopK(build(), k, b, avgDL, topK, filterSet, useBMW) + require.Lenf(t, got, len(want), "trial %d filtered bmw=%v len", trial, useBMW) + for i := range want { + require.InEpsilonf(t, want[i].score, got[i].score, 1e-9, + "trial %d filtered bmw=%v rank %d score", trial, useBMW, i) + tied := (i > 0 && wantPlus[i].score == wantPlus[i-1].score) || + (i+1 < len(wantPlus) && wantPlus[i].score == wantPlus[i+1].score) + if !tied { + require.Equalf(t, want[i].uid, got[i].uid, + "trial %d filtered bmw=%v rank %d uid", trial, useBMW, i) + } + } + } + } +} + +// TestWandTieHeavyDifferentialWindows (Q3) is a randomized differential over +// tie-heavy multi-term corpora: wandTopK (BMW on and off, varying topK, with and +// without a filterSet) must return exactly the first-topK window of the scoreAllDocs +// total order (score desc, uid asc) — same uids in the same order, bit-identical +// scores. wandTopK accumulates a document's score doc-major over cursors ordered by +// an unstable sort while scoreAllDocs accumulates term-major, so the corpus is built +// so that every per-term contribution is an exact dyadic float: k=1 and b=0 make +// bm25Score(idf, tf, ...) = idf*2*tf/(1+tf), which with TF in {1,3} and power-of-two +// IDFs yields contributions in {idf, 1.5*idf} — small multiples of 0.25 whose sums +// are exact in ANY addition order. Score ties are therefore EXACT and plentiful, and +// any window mismatch is a genuine heap/pruning defect, not float-grouping noise. +// (Whether last-ulp grouping differences on non-dyadic ties can flip page boundaries +// is an open semantics question covered by the integration-level pagination tests, +// not decidable in a deterministic unit test.) +func TestWandTieHeavyDifferentialWindows(t *testing.T) { + rng := rand.New(rand.NewSource(1234)) + k, b, avgDL := 1.0, 0.0, 10.0 + idfChoices := []float64{0.5, 1.0, 2.0} + tfChoices := []uint32{1, 3} + + for trial := 0; trial < 150; trial++ { + numTerms := 2 + rng.Intn(3) // 2..4 terms + numDocs := 60 + rng.Intn(440) // up to ~300 postings/term: spans >1 block + termPostings := make([][]posting.BM25Posting, numTerms) + idfs := make([]float64, numTerms) + for ti := 0; ti < numTerms; ti++ { + var ps []posting.BM25Posting + for i := 0; i < numDocs; i++ { + if rng.Intn(10) >= 6 { // ~60% of docs match each term: heavy overlap + continue + } + ps = append(ps, posting.BM25Posting{ + Uid: uint64(i + 1), // built ascending, as real posting lists are + TF: tfChoices[rng.Intn(len(tfChoices))], + DocLen: uint32(1 + rng.Intn(30)), // inert with b=0; exercises block bounds + }) + } + termPostings[ti] = ps + idfs[ti] = idfChoices[rng.Intn(len(idfChoices))] + } + + build := func() []*termCursor { + cs := make([]*termCursor, 0, numTerms) + for ti, ps := range termPostings { + if len(ps) == 0 { + continue + } + cs = append(cs, newTermCursor(ps, idfs[ti], k, b, avgDL)) + } + return cs + } + + // Random filter subset, exercised alongside the unfiltered run. + filterSet := map[uint64]struct{}{} + for i := 0; i < numDocs; i++ { + if rng.Intn(2) == 0 { + filterSet[uint64(i+1)] = struct{}{} + } + } + + for _, fs := range []map[uint64]struct{}{nil, filterSet} { + full := scoreAllDocs(build(), k, b, avgDL, fs) + for _, topK := range []int{1, 2, 3, 7, 20, len(full), len(full) + 5} { + if topK <= 0 { + continue + } + want := full + if topK < len(want) { + want = want[:topK] + } + for _, useBMW := range []bool{false, true} { + got := wandTopK(build(), k, b, avgDL, topK, fs, useBMW) + require.Lenf(t, got, len(want), + "trial %d filtered=%v bmw=%v topK=%d len", trial, fs != nil, useBMW, topK) + for i := range want { + require.Equalf(t, want[i].uid, got[i].uid, + "trial %d filtered=%v bmw=%v topK=%d rank %d uid", + trial, fs != nil, useBMW, topK, i) + require.Equalf(t, + math.Float64bits(want[i].score), math.Float64bits(got[i].score), + "trial %d filtered=%v bmw=%v topK=%d rank %d score bits", + trial, fs != nil, useBMW, topK, i) + } + } + } + } + } +} + +// TestWandFreshBoundsDifferential (Q14) pins that WAND/BMW block upper bounds are +// recomputed from the live postings on every search — nothing is persisted or cached +// across queries, so mutations between searches can never leave stale bounds that +// prune a new top scorer. Tie-free corpus; after each mutation the cursors are +// rebuilt (as wandSearch does per query) and BMW must stay byte-identical to the +// exhaustive path. Guards against a future change that caches block maxima. +func TestWandFreshBoundsDifferential(t *testing.T) { + k, b, avgDL := 1.2, 0.75, 15.0 + idfs := []float64{1.7, 0.9} + const topK = 10 + + // Two overlapping terms, term0 spanning multiple wandBlockSize blocks. Distinct + // (TF, DocLen) mixes make summed scores strictly distinct (asserted below), so + // the top-k uid list is unambiguous. + term0 := make([]posting.BM25Posting, 0, 300) + term1 := make([]posting.BM25Posting, 0, 150) + for i := 0; i < 300; i++ { + uid := uint64(i + 1) + term0 = append(term0, posting.BM25Posting{ + Uid: uid, TF: uint32(1 + (i*7)%13), DocLen: uint32(5 + i), + }) + if i%2 == 0 { + term1 = append(term1, posting.BM25Posting{ + Uid: uid, TF: uint32(1 + (i*5)%9), DocLen: uint32(5 + i), + }) + } + } + termPostings := [][]posting.BM25Posting{term0, term1} + + build := func() []*termCursor { + cs := make([]*termCursor, 0, len(termPostings)) + for ti, ps := range termPostings { + if len(ps) == 0 { + continue + } + cs = append(cs, newTermCursor(ps, idfs[ti], k, b, avgDL)) + } + return cs + } + + // checkDifferential rebuilds cursors from the CURRENT postings (fresh bounds, + // exactly as wandSearch does per query) and asserts WAND and BMW byte-identical + // to exhaustive scoring. Returns the full ranking for the next mutation step. + checkDifferential := func(phase string) []scoredDoc { + full := scoreAllDocs(build(), k, b, avgDL, nil) + require.NotEmpty(t, full, "%s: corpus must not be empty", phase) + for i := 1; i < len(full); i++ { + require.NotEqualf(t, full[i-1].score, full[i].score, + "%s: corpus must stay tie-free (ranks %d/%d)", phase, i-1, i) + } + want := full + if topK < len(want) { + want = want[:topK] + } + for _, useBMW := range []bool{false, true} { + got := wandTopK(build(), k, b, avgDL, topK, nil, useBMW) + require.Lenf(t, got, len(want), "%s bmw=%v len", phase, useBMW) + for i := range want { + require.Equalf(t, want[i].uid, got[i].uid, "%s bmw=%v rank %d uid", + phase, useBMW, i) + require.Equalf(t, + math.Float64bits(want[i].score), math.Float64bits(got[i].score), + "%s bmw=%v rank %d score bits", phase, useBMW, i) + } + } + return full + } + + full := checkDifferential("baseline") + + // Mutation 1: delete the top-scored document (the doc setting its block's max + // bound) from every term. Recomputed bounds must reflect the deletion — the uid + // must vanish and the differential must still hold exactly. + topUID := full[0].uid + for ti, ps := range termPostings { + kept := make([]posting.BM25Posting, 0, len(ps)) + for _, p := range ps { + if p.Uid != topUID { + kept = append(kept, p) + } + } + termPostings[ti] = kept + } + full = checkDifferential("after-delete-top") + for _, d := range full { + require.NotEqual(t, topUID, d.uid, "deleted doc must not be returned") + } + + // Mutation 2: promote a low-ranked doc to the top scorer by boosting its TF far + // above every block's previous max (and shrinking DocLen) in BOTH terms — BM25 + // saturates at idf*(k+1) per term, so only a doc matching both terms can beat + // docs that match both. Stale cached bounds would prune it; fresh bounds must + // rank it first. + var promoted uint64 + for i := len(full) - 1; i >= 0; i-- { + if (full[i].uid-1)%2 == 0 { // uid = i+1 with even i: present in term1 too + promoted = full[i].uid + break + } + } + require.NotZero(t, promoted, "corpus must contain a low-ranked doc in both terms") + for ti := range termPostings { + for i := range termPostings[ti] { + if termPostings[ti][i].Uid == promoted { + termPostings[ti][i].TF = 200 + termPostings[ti][i].DocLen = 1 + } + } + } + full = checkDifferential("after-promote-bottom") + require.Equal(t, promoted, full[0].uid, + "promoted doc must be the new top scorer under freshly computed bounds") +} + +// TestWandMatchesBruteForce checks that WAND and Block-Max WAND return exactly the +// same top-k documents and scores as exhaustive scoring, across many randomized +// posting lists. This is the core correctness guarantee: pruning must never change +// the result, only the work done. +func TestWandMatchesBruteForce(t *testing.T) { + rng := rand.New(rand.NewSource(42)) + k, b, avgDL := 1.2, 0.75, 12.0 + + for trial := 0; trial < 200; trial++ { + numTerms := 1 + rng.Intn(4) + termPostings := make([][]posting.BM25Posting, numTerms) + idfs := make([]float64, numTerms) + for ti := 0; ti < numTerms; ti++ { + n := rng.Intn(400) // spans multiple wandBlockSize blocks + seen := map[uint64]bool{} + var ps []posting.BM25Posting + for j := 0; j < n; j++ { + uid := uint64(1 + rng.Intn(500)) + if seen[uid] { + continue + } + seen[uid] = true + ps = append(ps, posting.BM25Posting{ + Uid: uid, + TF: uint32(1 + rng.Intn(10)), + DocLen: uint32(1 + rng.Intn(30)), + }) + } + sort.Slice(ps, func(i, j int) bool { return ps[i].Uid < ps[j].Uid }) + termPostings[ti] = ps + // Vary IDF per term so different terms carry different weight. + idfs[ti] = 0.5 + rng.Float64()*2 + } + + topK := 1 + rng.Intn(10) + want := bruteForceTopK(termPostings, idfs, k, b, avgDL, topK) + // One extra result lets us detect a tie between the cutoff rank and the + // first excluded document (a boundary tie outside the top-k window). + wantPlus := bruteForceTopK(termPostings, idfs, k, b, avgDL, topK+1) + + build := func() []*termCursor { + cs := make([]*termCursor, 0, numTerms) + for ti, ps := range termPostings { + if len(ps) == 0 { + continue + } + cs = append(cs, newTermCursor(ps, idfs[ti], k, b, avgDL)) + } + return cs + } + + for _, useBMW := range []bool{false, true} { + got := wandTopK(build(), k, b, avgDL, topK, nil, useBMW) + require.Lenf(t, got, len(want), "trial %d bmw=%v len", trial, useBMW) + for i := range want { + // The score at each rank must match exactly: WAND/BMW pruning must + // never change which scores make the top-k, only the work done. + require.InEpsilonf(t, want[i].score, got[i].score, 1e-9, + "trial %d bmw=%v rank %d score", trial, useBMW, i) + // The uid is only guaranteed when this rank's score is not tied with + // a neighbor (including the first excluded doc); tied-boundary docs + // are interchangeable in the ranking. + tied := (i > 0 && wantPlus[i].score == wantPlus[i-1].score) || + (i+1 < len(wantPlus) && wantPlus[i].score == wantPlus[i+1].score) + if !tied { + require.Equalf(t, want[i].uid, got[i].uid, + "trial %d bmw=%v rank %d uid", trial, useBMW, i) + } + } + } + } +} diff --git a/worker/mutation.go b/worker/mutation.go index 2aaf7b2f4ee..1030cde2d50 100644 --- a/worker/mutation.go +++ b/worker/mutation.go @@ -409,6 +409,19 @@ func checkSchema(s *pb.SchemaUpdate) error { x.ParseAttr(s.Predicate)) } + // BM25 scores a single document (one value) per UID: per-document length and + // corpus statistics are not well-defined for a list predicate, and the bucketed + // stats maintenance relies on conflict detection that a list predicate's + // value-dependent conflict key would not provide. Reject @index(bm25) on lists. + if s.List { + for _, tokenizer := range s.Tokenizer { + if tokenizer == "bm25" { + return errors.Errorf("Tokenizer 'bm25' cannot be applied to list predicate: %s", + x.ParseAttr(s.Predicate)) + } + } + } + // If schema update has upsert directive, it should have index directive. if s.Upsert && len(s.Tokenizer) == 0 && !s.Unique { return errors.Errorf("Index tokenizer is mandatory for: [%s] when specifying @upsert directive", diff --git a/worker/mutation_integration_test.go b/worker/mutation_integration_test.go index 99a2a1eed01..f1f4b81695b 100644 --- a/worker/mutation_integration_test.go +++ b/worker/mutation_integration_test.go @@ -93,6 +93,25 @@ func TestCheckSchema(t *testing.T) { } require.NoError(t, checkSchema(s1)) + // bm25 on a scalar string predicate is allowed. + s1 = &pb.SchemaUpdate{ + Predicate: x.AttrInRootNamespace("bio"), + ValueType: pb.Posting_STRING, + Directive: pb.SchemaUpdate_INDEX, + Tokenizer: []string{"bm25"}, + } + require.NoError(t, checkSchema(s1)) + + // bm25 on a list predicate is rejected. + s1 = &pb.SchemaUpdate{ + Predicate: x.AttrInRootNamespace("tags"), + ValueType: pb.Posting_STRING, + Directive: pb.SchemaUpdate_INDEX, + Tokenizer: []string{"bm25"}, + List: true, + } + require.Error(t, checkSchema(s1)) + s1 = &pb.SchemaUpdate{ Predicate: x.AttrInRootNamespace("friend"), ValueType: pb.Posting_UID, diff --git a/worker/task.go b/worker/task.go index 409ec3f0fc4..99711da1b3d 100644 --- a/worker/task.go +++ b/worker/task.go @@ -7,6 +7,7 @@ package worker import ( "context" + "encoding/binary" "fmt" "math" "sort" @@ -224,6 +225,7 @@ const ( customIndexFn matchFn similarToFn + bm25SearchFn standardFn = 100 ) @@ -266,6 +268,8 @@ func parseFuncTypeHelper(name string) (FuncType, string) { return uidInFn, f case "similar_to": return similarToFn, f + case "bm25": + return bm25SearchFn, f case "anyof", "allof": return customIndexFn, f case "match": @@ -292,6 +296,8 @@ func needsIndex(fnType FuncType, uidList *pb.List) bool { return true case similarToFn: return true + case bm25SearchFn: + return true } return false } @@ -314,7 +320,7 @@ type funcArgs struct { // The function tells us whether we want to fetch value posting lists or uid posting lists. func (srcFn *functionContext) needsValuePostings(typ types.TypeID) (bool, error) { switch srcFn.fnType { - case aggregatorFn, passwordFn, similarToFn: + case aggregatorFn, passwordFn, similarToFn, bm25SearchFn: return true, nil case compareAttrFn: if len(srcFn.tokens) > 0 { @@ -351,16 +357,27 @@ func (qs *queryState) handleValuePostings(ctx context.Context, args funcArgs) er attribute.String("srcFn", x.SafeUTF8(fmt.Sprintf("%+v", args.srcFn))))) switch srcFn.fnType { - case notAFunction, aggregatorFn, passwordFn, compareAttrFn, similarToFn: + case notAFunction, aggregatorFn, passwordFn, compareAttrFn, similarToFn, bm25SearchFn: default: return errors.Errorf("Unhandled function in handleValuePostings: %s", srcFn.fname) } + if srcFn.fnType == bm25SearchFn { + return qs.handleBM25Search(ctx, args) + } + if srcFn.fnType == similarToFn { numNeighbors, err := strconv.ParseInt(q.SrcFunc.Args[0], 10, 32) if err != nil { return fmt.Errorf("invalid value for number of neighbors: %s", q.SrcFunc.Args[0]) } + // HNSW panics on expectedNeighbors <= 0 (searchLayerResult truncates its + // neighbor slice to zero length and dereferences neighbors[0]); reject the + // value before it reaches the index rather than crashing the alpha. + if numNeighbors <= 0 { + return fmt.Errorf("similar_to requires a positive number of neighbors, got %d", + numNeighbors) + } cspec, err := pickFactoryCreateSpec(ctx, args.q.Attr) if err != nil { return err @@ -375,6 +392,7 @@ func (qs *queryState) handleValuePostings(ctx context.Context, args funcArgs) er return err } var nnUids []uint64 + var nnScores []float64 // Build optional search options if provided filter := index.AcceptAll[float32] opts := index.VectorIndexOptions[float32]{Filter: filter} @@ -385,27 +403,63 @@ func (qs *queryState) handleValuePostings(ctx context.Context, args funcArgs) er opts.DistanceThreshold = srcFn.vsDistanceThreshold } hasOptions := opts.EfOverride > 0 || opts.DistanceThreshold != nil - if o, ok := indexer.(index.OptionalSearchOptions[float32]); ok && hasOptions { - if srcFn.vectorInfo != nil { - nnUids, err = o.SearchWithOptions(ctx, qc, srcFn.vectorInfo, int(numNeighbors), opts) - } else { - nnUids, err = o.SearchWithUidAndOptions(ctx, qc, srcFn.vectorUid, int(numNeighbors), opts) + // Use the scored search path so the per-uid similarity score can be bound to a + // value variable (powering native hybrid search / fuse()). The scored variants + // mirror their unscored counterparts exactly — SearchScored/SearchWithUidScored + // preserve the plain-query neighbor selection, and the *Options* variants are + // used only when the query supplies ef/distance-threshold — so adding scoring + // does not change which neighbors existing queries return. Indexes that don't + // implement scoring fall back to the unscored path (no scores surfaced). + if so, ok := indexer.(index.ScoredSearchOptions[float32]); ok { + switch { + case hasOptions && srcFn.vectorInfo != nil: + nnUids, nnScores, err = so.SearchWithOptionsScored(ctx, qc, srcFn.vectorInfo, int(numNeighbors), opts) + case hasOptions: + nnUids, nnScores, err = so.SearchWithUidAndOptionsScored(ctx, qc, srcFn.vectorUid, int(numNeighbors), opts) + case srcFn.vectorInfo != nil: + nnUids, nnScores, err = so.SearchScored(ctx, qc, srcFn.vectorInfo, int(numNeighbors), index.AcceptAll[float32]) + default: + nnUids, nnScores, err = so.SearchWithUidScored(ctx, qc, srcFn.vectorUid, int(numNeighbors), index.AcceptAll[float32]) } + } else if srcFn.vectorInfo != nil { + nnUids, err = indexer.Search(ctx, qc, srcFn.vectorInfo, + int(numNeighbors), index.AcceptAll[float32]) } else { - if srcFn.vectorInfo != nil { - nnUids, err = indexer.Search(ctx, qc, srcFn.vectorInfo, - int(numNeighbors), index.AcceptAll[float32]) - } else { - nnUids, err = indexer.SearchWithUid(ctx, qc, srcFn.vectorUid, - int(numNeighbors), index.AcceptAll[float32]) - } + nnUids, err = indexer.SearchWithUid(ctx, qc, srcFn.vectorUid, + int(numNeighbors), index.AcceptAll[float32]) } if err != nil && !strings.Contains(err.Error(), hnsw.EmptyHNSWTreeError+": "+badger.ErrKeyNotFound.Error()) { return err } - sort.Slice(nnUids, func(i, j int) bool { return nnUids[i] < nnUids[j] }) - args.out.UidMatrix = append(args.out.UidMatrix, &pb.List{Uids: nnUids}) + + // Emit uids ascending (required by the query pipeline) with positionally + // aligned similarity scores in ValueMatrix. The query layer binds these to a + // value variable so callers can order by and project the score via val(), and + // so vector results can serve as a fuse() channel. Scores are higher-is-better. + order := make([]int, len(nnUids)) + for i := range order { + order[i] = i + } + sort.Slice(order, func(i, j int) bool { return nnUids[order[i]] < nnUids[order[j]] }) + sortedUids := make([]uint64, len(nnUids)) + for i, idx := range order { + sortedUids[i] = nnUids[idx] + } + args.out.UidMatrix = append(args.out.UidMatrix, &pb.List{Uids: sortedUids}) + + if len(nnScores) == len(nnUids) && len(nnScores) > 0 { + scoreBuf := make([]byte, len(nnUids)*8) + scoreValues := make([]*pb.ValueList, len(nnUids)) + for i, idx := range order { + off := i * 8 + binary.LittleEndian.PutUint64(scoreBuf[off:off+8], math.Float64bits(nnScores[idx])) + scoreValues[i] = &pb.ValueList{ + Values: []*pb.TaskValue{{Val: scoreBuf[off : off+8 : off+8], ValType: pb.Posting_FLOAT}}, + } + } + args.out.ValueMatrix = append(args.out.ValueMatrix, scoreValues...) + } return nil } @@ -1219,6 +1273,115 @@ func needsStringFiltering(srcFn *functionContext, langs []string, attr string) b srcFn.fnType == customIndexFn || srcFn.fnType == ngramFn) } +func (qs *queryState) handleBM25Search(ctx context.Context, args funcArgs) error { + q := args.q + attr := q.Attr + + // 1. Parse args: query text, optional k (default 1.2), b (default 0.75). + if len(q.SrcFunc.Args) < 1 { + return errors.Errorf("bm25 requires at least 1 argument (query text)") + } + queryText := q.SrcFunc.Args[0] + k := 1.2 + b := 0.75 + if len(q.SrcFunc.Args) >= 2 { + var err error + k, err = strconv.ParseFloat(q.SrcFunc.Args[1], 64) + if err != nil { + return errors.Errorf("bm25: invalid k parameter: %s", q.SrcFunc.Args[1]) + } + } + if len(q.SrcFunc.Args) >= 3 { + var err error + b, err = strconv.ParseFloat(q.SrcFunc.Args[2], 64) + if err != nil { + return errors.Errorf("bm25: invalid b parameter: %s", q.SrcFunc.Args[2]) + } + } + if math.IsNaN(k) || math.IsInf(k, 0) || k <= 0 { + return errors.Errorf("bm25: k must be a positive finite number, got %v", k) + } + if math.IsNaN(b) || math.IsInf(b, 0) || b < 0 || b > 1 { + return errors.Errorf("bm25: b must be between 0 and 1, got %v", b) + } + + // 2. Tokenize query (deduplicated) using the fulltext pipeline. The returned + // tokens already carry the BM25 tokenizer identifier byte. + lang := langForFunc(q.Langs) + queryTokens, err := tok.GetBM25QueryTokens([]string{queryText}, lang) + if err != nil { + return err + } + if len(queryTokens) == 0 { + args.out.UidMatrix = append(args.out.UidMatrix, &pb.List{}) + return nil + } + + // 3. Read bucketed corpus stats and derive N and the average document length. + docCount, totalTerms, err := posting.ReadBM25Stats(qs.cache.Get, attr, q.ReadTs) + if err != nil { + return err + } + if docCount == 0 || totalTerms == 0 { + args.out.UidMatrix = append(args.out.UidMatrix, &pb.List{}) + return nil + } + avgDL := float64(totalTerms) / float64(docCount) + N := float64(docCount) + + // Build a filter set if bm25 is used as a filter (@filter(bm25(...))). + var filterSet map[uint64]struct{} + if q.UidList != nil && len(q.UidList.Uids) > 0 { + filterSet = make(map[uint64]struct{}, len(q.UidList.Uids)) + for _, uid := range q.UidList.Uids { + filterSet[uid] = struct{}{} + } + } + + // 4. Use WAND top-k early termination whenever a first limit is set, retaining + // first+offset documents so the offset can be dropped afterward. This bounds work + // and memory to the requested window instead of scoring (and materializing) every + // matching document just because an offset was supplied. With no first limit, + // topK is 0 and every match is scored (the caller explicitly asked for all results). + topK := bm25TopK(int(q.First), int(q.Offset)) + + // 5. Run WAND / Block-Max WAND over the standard posting lists. + results, err := wandSearch(qs.cache.Get, attr, q.ReadTs, queryTokens, k, b, avgDL, N, + topK, filterSet, true) + if err != nil { + return err + } + + // 6. Apply the first/offset window over the score-descending results. wandSearch + // returns results sorted by score (descending), whether or not it top-k'd them, so + // the same slice is correct for both the top-k and score-all paths. + results = bm25PaginateScored(results, int(q.First), int(q.Offset)) + + // 7. Emit UIDs ascending (required by the query pipeline) with positionally- + // aligned scores in ValueMatrix; the query layer binds these to a value + // variable so callers order by and project the score via val(). + sort.Slice(results, func(i, j int) bool { return results[i].uid < results[j].uid }) + uids := make([]uint64, len(results)) + for i, r := range results { + uids[i] = r.uid + } + args.out.UidMatrix = append(args.out.UidMatrix, &pb.List{Uids: uids}) + + scoreBuf := make([]byte, len(results)*8) + scoreValues := make([]*pb.ValueList, len(results)) + for i, r := range results { + off := i * 8 + binary.LittleEndian.PutUint64(scoreBuf[off:off+8], math.Float64bits(r.score)) + // Three-index slice caps capacity at 8 so a downstream append can't corrupt + // adjacent scores in the shared backing array. + scoreValues[i] = &pb.ValueList{ + Values: []*pb.TaskValue{{Val: scoreBuf[off : off+8 : off+8], ValType: pb.Posting_FLOAT}}, + } + } + args.out.ValueMatrix = append(args.out.ValueMatrix, scoreValues...) + return nil +} + func (qs *queryState) handleCompareScalarFunction(ctx context.Context, arg funcArgs) error { attr := arg.q.Attr if ok := schema.State().HasCount(ctx, attr); !ok { @@ -2167,6 +2330,18 @@ func parseSrcFn(ctx context.Context, q *pb.Query) (*functionContext, error) { return nil, err } checkRoot(q, fc) + case bm25SearchFn: + // bm25(pred, "query text") or bm25(pred, "query text", "k", "b") + if len(q.SrcFunc.Args) < 1 || len(q.SrcFunc.Args) > 3 { + return nil, errors.Errorf("Function 'bm25' requires 1-3 arguments (query [, k, b]), but got %d", + len(q.SrcFunc.Args)) + } + required, found := verifyStringIndex(ctx, attr, fnType) + if !found { + return nil, errors.Errorf("Attribute %s is not indexed with type %s", x.ParseAttr(attr), + required) + } + checkRoot(q, fc) case similarToFn: // similar_to accepts 2 mandatory args: k, vector_or_uid followed by optional key:value pairs // Example: similar_to(vpred, 3, $vec, ef: 64, distance_threshold: 0.5) @@ -2178,6 +2353,16 @@ func parseSrcFn(ctx context.Context, q *pb.Query) (*functionContext, error) { if err != nil { return nil, err } + // Reject non-finite query-vector components up front: a NaN poisons every + // distance comparison inside HNSW (undefined ordering, arbitrary neighbors) + // and would surface as silently wrong similarity scores. + for _, f := range fc.vectorInfo { + if math.IsNaN(float64(f)) || math.IsInf(float64(f), 0) { + return nil, errors.Errorf( + "Function '%s' requires a finite query vector; got a NaN/Inf component", + q.SrcFunc.Name) + } + } if len(q.SrcFunc.Args) > 2 { if err := parseSimilarToOptions(q.SrcFunc.Args[2:], fc); err != nil { return nil, err diff --git a/worker/tokens.go b/worker/tokens.go index 2740d29f447..b8c85a22816 100644 --- a/worker/tokens.go +++ b/worker/tokens.go @@ -25,6 +25,8 @@ func verifyStringIndex(ctx context.Context, attr string, funcType FuncType) (str requiredTokenizer = tok.NGramTokenizer{} case fullTextSearchFn: requiredTokenizer = tok.FullTextTokenizer{} + case bm25SearchFn: + requiredTokenizer = tok.BM25Tokenizer{} case matchFn: requiredTokenizer = tok.TrigramTokenizer{} default: @@ -65,6 +67,9 @@ func getStringTokens(funcArgs []string, lang string, funcType FuncType, query bo if funcType == fullTextSearchFn { return tok.GetFullTextTokens(funcArgs, lang) } + if funcType == bm25SearchFn { + return tok.GetBM25QueryTokens(funcArgs, lang) + } if funcType == ngramFn { if query { return tok.GetNGramQueryTokens(funcArgs, lang) diff --git a/x/keys.go b/x/keys.go index 1607f20b2cc..1883f026c14 100644 --- a/x/keys.go +++ b/x/keys.go @@ -292,6 +292,51 @@ func CountKey(attr string, count uint32, reverse bool) []byte { return buf } +// BM25IndexKey generates the index key for a BM25 term posting list. The +// encodedToken already carries the BM25 tokenizer identifier byte, so BM25 term +// postings live at the same standard index key as every other tokenizer — +// IndexKey(attr, identifier || term) — and inherit rollup, splits, backup, and +// index-rebuild handling for free. This is a thin alias of IndexKey so the index +// write path and the query read path share one definition. +func BM25IndexKey(attr string, encodedToken string) []byte { + return IndexKey(attr, encodedToken) +} + +// bm25StatsPrefix namespaces the BM25 corpus-statistics keys. These hold the +// document count and total term count (used to derive the average document +// length); they are auxiliary metadata, not term postings, so they use a reserved +// token that cannot collide with any stemmed BM25 term. +const bm25StatsPrefix = "\x00_bm25stats_" + +// BM25StatsKey generates the key for one bucket of BM25 corpus statistics. Stats +// are sharded across buckets (keyed by uid%numBuckets) to spread write contention. +func BM25StatsKey(attr string, bucket int) []byte { + var buf [2]byte + binary.BigEndian.PutUint16(buf[:], uint16(bucket)) + return IndexKey(attr, bm25StatsPrefix+string(buf[:])) +} + +// IsBM25Stats reports whether p is a BM25 corpus-statistics bucket key (the +// reserved-token index keys written by BM25StatsKey). The bulk reducer uses this +// to fold per-document stats postings into the single aggregate bucket posting. +func (p ParsedKey) IsBM25Stats() bool { + return p.IsIndex() && strings.HasPrefix(p.Term, bm25StatsPrefix) +} + +// BM25StatsPrefix returns the key prefix covering every BM25 corpus-statistics +// bucket for attr. Stats live under a reserved token prefix that is distinct from +// the IdentBM25 term-posting prefix, so dropping/rebuilding the BM25 index must +// delete this prefix too — otherwise the stale stats survive a drop and +// double-count when the index is rebuilt on top of them. The split argument +// selects the ByteSplit-prefixed variant used by multi-part lists. +func BM25StatsPrefix(attr string, split bool) []byte { + prefix := IndexKey(attr, bm25StatsPrefix) + if split { + prefix[0] = ByteSplit + } + return prefix +} + // ParsedKey represents a key that has been parsed into its multiple attributes. type ParsedKey struct { Attr string