fix(bm25): address Gemini/GPT-5 code review findings

- Add DecodeCount() to bm25enc for O(1) entry count reads without
  full decode, preventing OOM on legacy migration with large posting
  lists (e.g., common terms with millions of entries)
- Use DecodeCount in WAND search legacy DF calculation path
- Fix integer overflow in DecodeDir bounds check by using uint64
  arithmetic (prevents panic on corrupted data with MaxUint32 count)
- Pre-allocate shared score buffer in handleBM25Search with
  three-index slices to prevent accidental append corruption
- Document bm25Writes concurrency model and limitations

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

Author: shaunpatterson

posting/bm25block/bm25block.go

@@ -77,7 +77,8 @@ func DecodeDir(data []byte) *Dir {
 	}
 	count := binary.BigEndian.Uint32(data[0:4])
 	nextID := binary.BigEndian.Uint32(data[4:8])
-	if int(count)*dirEntrySize+dirHeaderSize > len(data) {
+	// Use uint64 arithmetic to prevent integer overflow on corrupted data.
+	if uint64(count)*dirEntrySize+dirHeaderSize > uint64(len(data)) {
 		return &Dir{NextID: nextID}
 	}
 	blocks := make([]BlockMeta, count)

posting/bm25block/bm25block_test.go

@@ -6,6 +6,7 @@
 package bm25block
 
 import (
+	"encoding/binary"
 	"math"
 	"testing"
 
@@ -39,6 +40,18 @@ func TestDirRoundtripEmpty(t *testing.T) {
 	require.Empty(t, got.Blocks)
 }
 
+func TestDecodeDirCorruptedLargeCount(t *testing.T) {
+	// A corrupted blob with a massive count should not panic due to integer overflow.
+	// count = MaxUint32, nextID = 0, followed by only 8 bytes of data.
+	data := make([]byte, 16)
+	binary.BigEndian.PutUint32(data[0:4], 0xFFFFFFFF) // count = MaxUint32
+	binary.BigEndian.PutUint32(data[4:8], 0)          // nextID = 0
+	got := DecodeDir(data)
+	// Should return an empty Dir (with nextID preserved) rather than panicking.
+	require.Empty(t, got.Blocks)
+	require.Equal(t, uint32(0), got.NextID)
+}
+
 func TestDirRoundtripSingle(t *testing.T) {
 	dir := &Dir{
 		NextID: 1,

posting/bm25enc/bm25enc.go

@@ -130,6 +130,17 @@ func UIDs(entries []Entry) []uint64 {
 	return uids
 }
 
+// DecodeCount reads just the entry count from the header of an encoded blob
+// without decoding any entries. This is O(1) and avoids allocating a full
+// []Entry slice, which matters for large posting lists (e.g., common terms
+// during legacy format migration).
+func DecodeCount(data []byte) uint32 {
+	if len(data) < 4 {
+		return 0
+	}
+	return binary.BigEndian.Uint32(data[:4])
+}
+
 // EncodeStats encodes BM25 corpus statistics (docCount, totalTerms) as 16 bytes.
 func EncodeStats(docCount, totalTerms uint64) []byte {
 	buf := make([]byte, 16)

posting/bm25enc/bm25enc_test.go

@@ -92,6 +92,37 @@ func TestUIDs(t *testing.T) {
 	require.Equal(t, []uint64{1, 5, 100}, UIDs(entries))
 }
 
+func TestDecodeCount(t *testing.T) {
+	// Normal case: count matches actual entries.
+	entries := []Entry{
+		{UID: 1, Value: 3},
+		{UID: 5, Value: 1},
+		{UID: 100, Value: 7},
+	}
+	data := Encode(entries)
+	require.Equal(t, uint32(3), DecodeCount(data))
+
+	// Empty/nil input.
+	require.Equal(t, uint32(0), DecodeCount(nil))
+	require.Equal(t, uint32(0), DecodeCount([]byte{}))
+	require.Equal(t, uint32(0), DecodeCount([]byte{1, 2, 3}))
+
+	// Zero count.
+	require.Equal(t, uint32(0), DecodeCount([]byte{0, 0, 0, 0}))
+
+	// Single entry.
+	single := Encode([]Entry{{UID: 42, Value: 10}})
+	require.Equal(t, uint32(1), DecodeCount(single))
+
+	// Large count.
+	large := make([]Entry, 10000)
+	for i := range large {
+		large[i] = Entry{UID: uint64(i*3 + 1), Value: uint32(i % 100)}
+	}
+	data = Encode(large)
+	require.Equal(t, uint32(10000), DecodeCount(data))
+}
+
 func TestStatsRoundtrip(t *testing.T) {
 	data := EncodeStats(12345, 98765)
 	dc, tt := DecodeStats(data)

posting/lists.go

@@ -79,6 +79,18 @@ type LocalCache struct {
 
 	// bm25Writes buffers BM25 direct KV writes (key → encoded blob).
 	// These bypass the posting list infrastructure entirely.
+	//
+	// CONCURRENCY NOTE: BM25 blocks use full-value overwrites rather than
+	// posting list deltas. Within a single Dgraph transaction this is safe
+	// (each Txn has its own LocalCache). Across concurrent transactions,
+	// Dgraph's Raft-based mutation serialization prevents lost updates for
+	// the same predicate+UID pair. However, two transactions updating
+	// different UIDs that share a common term could theoretically race on
+	// the same term block. In practice this is mitigated by:
+	//   1. Dgraph serializes mutations through Raft proposals
+	//   2. Block splits keep contention surface small
+	// If higher write concurrency is needed, blocks should be integrated
+	// into the posting list delta mechanism.
 	bm25Writes map[string][]byte
 }
 

worker/bm25wand.go

@@ -447,11 +447,11 @@ func wandSearch(attr string, readTs uint64, queryTokens []string,
 				df += uint64(bm.Count)
 			}
 		} else {
-			// Legacy fallback: read the monolithic blob to get df.
+			// Legacy fallback: read just the count header to get df.
+			// Avoids decoding the full posting list (which could be huge for common terms).
 			legacyKey := x.BM25IndexKey(attr, token)
 			legacyBlob := posting.ReadBM25BlobAt(legacyKey, readTs)
-			legacyEntries := bm25enc.Decode(legacyBlob)
-			df = uint64(len(legacyEntries))
+			df = uint64(bm25enc.DecodeCount(legacyBlob))
 		}
 		if df == 0 {
 			continue

worker/task.go

@@ -1318,19 +1318,27 @@ func (qs *queryState) handleBM25Search(ctx context.Context, args funcArgs) error
 
 	// 7. Build output: UIDs sorted ascending (required by query pipeline)
 	// and ValueMatrix with aligned scores (for bm25_score pseudo-predicate).
+	// We use a single pre-allocated buffer for all score encodings to reduce
+	// per-result heap allocations.
 	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})
 
+	// Encode scores into ValueMatrix. Each entry in ValueMatrix corresponds
+	// positionally to a UID in UidMatrix[0], enabling the bm25_score
+	// pseudo-predicate in query.go to map UIDs to scores.
+	scoreBuf := make([]byte, len(results)*8)
 	scoreValues := make([]*pb.ValueList, len(results))
 	for i, r := range results {
-		buf := make([]byte, 8)
-		binary.LittleEndian.PutUint64(buf, math.Float64bits(r.score))
+		off := i * 8
+		binary.LittleEndian.PutUint64(scoreBuf[off:off+8], math.Float64bits(r.score))
+		// Use three-index slice to cap capacity at 8, preventing any downstream
+		// append from corrupting adjacent scores in the shared backing array.
 		scoreValues[i] = &pb.ValueList{
-			Values: []*pb.TaskValue{{Val: buf, ValType: pb.Posting_ValType(pb.Posting_FLOAT)}},
+			Values: []*pb.TaskValue{{Val: scoreBuf[off : off+8 : off+8], ValType: pb.Posting_ValType(pb.Posting_FLOAT)}},
 		}
 	}
 	args.out.ValueMatrix = append(args.out.ValueMatrix, scoreValues...)