diff --git a/src/CodeIndex.Core/AssemblyInfo.cs b/src/CodeIndex.Core/AssemblyInfo.cs new file mode 100644 index 0000000..aacad5b --- /dev/null +++ b/src/CodeIndex.Core/AssemblyInfo.cs @@ -0,0 +1,11 @@ +using System.Runtime.CompilerServices; + +// Grants CodeIndex.Core.Tests access to this assembly's internal members. Kept to the single +// existing use case rather than opened up broadly: VectorSearcher.ComputeScores and +// VectorSearcher.SelectTopKFromScores are internal (not private) purely so the realistic-scale +// benchmark test can measure the bounded-heap selection strategy in isolation from the SIMD +// dot-product pass it is fused with in the public Search method. Re-deriving that selection +// logic as a second, test-local copy was rejected — VectorSearcherTests already documents a real +// bug where its tie-break rule was easy to get subtly wrong, so a hand-rolled duplicate risked +// silently comparing two different algorithms instead of testing the real one. +[assembly: InternalsVisibleTo("CodeIndex.Core.Tests")] diff --git a/src/CodeIndex.Core/Search/VectorSearcher.cs b/src/CodeIndex.Core/Search/VectorSearcher.cs index 6cc7d13..5e8bda5 100644 --- a/src/CodeIndex.Core/Search/VectorSearcher.cs +++ b/src/CodeIndex.Core/Search/VectorSearcher.cs @@ -61,24 +61,13 @@ public VectorSearcher(float[] vectors, int dimensions) /// and scores identically under Reciprocal Rank Fusion to a genuinely strong match elsewhere. /// /// - /// Selection of the top K uses a size-bounded min-heap () - /// rather than sorting all scores: that is O(N log K) instead of - /// O(N log N). At the project's measured scale (8735 x 1024, topK 20, Release, best of - /// several warmed-up runs) this measured roughly 1.6 ms against roughly 2.1 ms for a full - /// sort — the win is real but modest (about half a millisecond here), so this is not the - /// dominant cost of a search; the SIMD scoring loop below, which both approaches share, is. - /// The heap never holds more than entries, so a candidate only - /// survives a comparison against the current worst kept score, not against the whole - /// result set. - /// - /// The heap is keyed by (Score, Index) rather than by score alone: - /// does not guarantee which of two - /// equal-priority entries is treated as the root, so keying by score alone let eviction - /// silently drop either side of a tie depending on internal heap shape — breaking the - /// ascending-index tie-break at the selection boundary even though the final sort still - /// enforced it among survivors. defines "worse" as lower - /// score, and on a tied score, as the *higher* index, so eviction always drops the - /// higher-index element of a tie and the kept set is deterministic regardless of topK. + /// Scoring () and selection () + /// are two separate steps chained together here. Selection uses a size-bounded min-heap + /// rather than sorting every score: O(N log K) instead of O(N log N) — see + /// for why, and for the tie-break rule it enforces at the + /// selection boundary. In practice the SIMD scoring pass, which any selection strategy must + /// pay identically, dominates the cost of a search; selection over it is a comparatively + /// small and fast step (see 's own remarks for scale). /// public IReadOnlyList Search(ReadOnlySpan query, int topK, float minScore = float.NegativeInfinity) { @@ -105,21 +94,95 @@ public IReadOnlyList Search(ReadOnlySpan query, int topK, fl int take = Math.Min(topK, _count); - PriorityQueue heap = new(take, WorstFirstComparer.Instance); + float[] scores = ComputeScores(query); + + return SelectTopKFromScores(scores, take, minScore); + } + + /// + /// Scores every row against via cosine similarity (a plain dot + /// product — see the class remarks on why vectors are assumed unit-normalised) and returns + /// the raw, unranked scores, one per row, in row order. + /// + /// + /// Split out from so the dot-product pass — the dominant, shared cost + /// of a search, identical regardless of selection strategy — can be measured or reused on + /// its own, separately from . internal rather than + /// private purely so the realistic-scale benchmark test can time the two independently + /// instead of only ever measuring them bundled together (see + /// VectorSearcherTests.Search_At8735By1024RealisticScale_CompletesWithinAGenerousBound + /// for why that bundling made the benchmark's heap-vs-full-sort ratio noisy). + /// + /// + /// This split is not free, and the cost lands on the production path rather than the test: + /// the previous fused loop scored and offered each row to the heap in one pass and so never + /// held more than topK entries, whereas materialising every score first makes a + /// search allocate an array of elements. Memory per + /// search goes from O(topK) to O(N). At this project's measured scale that is 8735 floats, + /// about 35 KB of short-lived Gen0 garbage per query, against a query that spends roughly + /// 190 ms waiting for Ollama to embed the text and about 1.6 ms searching — immaterial, and + /// measured to be so rather than assumed. It would stop being immaterial on an index one or + /// two orders of magnitude larger, at which point fusing the two steps back together for the + /// production path (keeping them separate only for the benchmark) is the fix. + /// + /// + internal float[] ComputeScores(ReadOnlySpan query) + { + float[] scores = new float[_count]; for (int i = 0; i < _count; i++) { - // Unit-normalised vectors make the dot product the cosine similarity directly. // TensorPrimitives.Dot runs as a single SIMD block operation over the whole row, // never element by element. - float score = TensorPrimitives.Dot(_vectors.AsSpan(i * _dimensions, _dimensions), query); + scores[i] = TensorPrimitives.Dot(_vectors.AsSpan(i * _dimensions, _dimensions), query); + } + + return scores; + } + + /// + /// Selects the highest values from (row + /// index i.e. array position doubles as the row's ), excluding + /// any value below , using a size-bounded min-heap rather than + /// sorting every score: that is O(N log ) instead of O(N log N). The + /// heap never holds more than entries, so a candidate only survives a + /// comparison against the current worst kept score, not against the whole result set. + /// + /// + /// The heap is keyed by (Score, Index) rather than by score alone: + /// does not guarantee which of two + /// equal-priority entries is treated as the root, so keying by score alone let eviction + /// silently drop either side of a tie depending on internal heap shape — breaking the + /// ascending-index tie-break at the selection boundary even though the final sort still + /// enforced it among survivors. defines "worse" as lower + /// score, and on a tied score, as the *higher* index, so eviction always drops the + /// higher-index element of a tie and the kept set is deterministic regardless of . + /// + /// internal static (rather than a private instance detail of ) for + /// the same reason as : it lets a benchmark exercise this + /// selection strategy in isolation, on a fixed precomputed array, + /// instead of only ever timing it fused with the dot-product pass. Reusing this exact method + /// — rather than a second, test-local copy of the heap logic — matters because the + /// eviction tie-break above is easy to get subtly wrong (see the regression tests built from + /// the duplicate-symbol scenario); a hand-rolled copy in test code could drift from this + /// implementation and silently start comparing two different algorithms. + /// + internal static ScoredIndex[] SelectTopKFromScores(ReadOnlySpan scores, int take, float minScore) + { + PriorityQueue heap = new(take, WorstFirstComparer.Instance); + + for (int i = 0; i < scores.Length; i++) + { + float score = scores[i]; if (score < minScore) { // Below the relevance floor: excluded outright, never merely low priority. A row // this weak must not fill out the result set just because fewer than `take` rows // cleared the floor — an empty (or short) result honestly says "nothing here was - // relevant enough," which is the whole point of the floor (see the parameter doc). + // relevant enough," which is the whole point of the floor (see Search's parameter + // doc). continue; } diff --git a/tests/CodeIndex.Core.Tests/Search/VectorSearcherTests.cs b/tests/CodeIndex.Core.Tests/Search/VectorSearcherTests.cs index 39593f9..839c0af 100644 --- a/tests/CodeIndex.Core.Tests/Search/VectorSearcherTests.cs +++ b/tests/CodeIndex.Core.Tests/Search/VectorSearcherTests.cs @@ -243,13 +243,31 @@ public void Search_TopKFromALargerLimitIsAPrefixOfTopKFromASmallerLimit() } /// - /// Realistic-scale timing at 8735 chunks x 1024 dimensions — the measured production - /// scale for this project. Reports the elapsed time for both the bounded-heap search and a - /// reference full sort to the test log, and asserts the two return identical results plus - /// a bound generous enough to not flake on a slow CI run while still being tight enough to - /// catch a real regression (a naive full sort at this scale would already be within this - /// bound, so it is not a meaningless tautology either). + /// Realistic-scale timing at 8735 chunks x 1024 dimensions — the measured production scale + /// for this project. Reports the elapsed time for both selection strategies to the test log, + /// asserts the full pipeline (score + select) matches a reference full sort, and separately + /// asserts a bound on the heap/full-sort ratio tight enough to catch a real regression. /// + /// + /// This test originally timed searcher.Search(...) (score + select fused) against a + /// reference that recomputed the same ~8735 x 1024 dot products and then did a full sort — + /// i.e. it measured "dot products + heap" against "dot products + full sort". Both paths pay + /// the identical, dominant dot-product cost, so in theory it should cancel out of the ratio. + /// In practice it did not always cancel: on one CI run the ratio spiked to 2.112 against a + /// 1.5 ceiling (issue #32), immediately passing on an identical re-run, while the full-sort + /// side's absolute time barely moved from its usual local value — pointing at something + /// landing specifically inside the heap path's measured window on that run (most likely GC + /// or scheduler noise from concurrent test collections, given xUnit's default + /// parallel-by-collection execution), not a real difference between the two algorithms. + /// + /// Rather than widen the ceiling to paper over that, this version computes the ~8735 x 1024 + /// scores once via , outside every timed region, + /// and times only against a full sort of + /// that same fixed scores array. That is what the assertion has always claimed to compare — + /// see , which is now the direct counterpart to + /// rather than to the whole of + /// . + /// [Fact] public void Search_At8735By1024RealisticScale_CompletesWithinAGenerousBound() { @@ -264,8 +282,19 @@ public void Search_At8735By1024RealisticScale_CompletesWithinAGenerousBound() float[] query = CreateRandomUnitVectors(1, dimensions, seed: 99); VectorSearcher searcher = new(vectors, dimensions); - IReadOnlyList hits = []; - ScoredIndex[] reference = []; + // The ~8735 x 1024 SIMD dot-product pass: identical work for either selection strategy, + // and the dominant cost of an end-to-end search. Computed once, here, outside every + // timed region below, so the loop measures only the two selection strategies against + // each other instead of diluting (and, per issue #32, occasionally skewing) the ratio + // with several milliseconds of cost neither strategy can avoid or differ on. + float[] scores = searcher.ComputeScores(query); + + // Correctness: the full pipeline (score + select) still has to match a plain full sort + // of the same scores. + IReadOnlyList hits = searcher.Search(query, topK); + ScoredIndex[] reference = FullSortSelectTopK(scores, topK); + Assert.Equal(topK, hits.Count); + Assert.Equal(reference, hits); // Tiered JIT compilation needs more than one call to reach steady-state optimised // code, and a single measured call (as this test originally took) can land mid-tier @@ -274,8 +303,8 @@ public void Search_At8735By1024RealisticScale_CompletesWithinAGenerousBound() // to see through GC pauses and scheduler noise in a micro-benchmark like this one. for (int i = 0; i < warmupRuns; i++) { - hits = searcher.Search(query, topK); - reference = FullSortReferenceSearch(vectors, dimensions, query, topK); + _ = VectorSearcher.SelectTopKFromScores(scores, topK, float.NegativeInfinity); + _ = FullSortSelectTopK(scores, topK); } double heapMs = double.MaxValue; @@ -284,52 +313,62 @@ public void Search_At8735By1024RealisticScale_CompletesWithinAGenerousBound() for (int i = 0; i < measuredRuns; i++) { Stopwatch heapStopwatch = Stopwatch.StartNew(); - hits = searcher.Search(query, topK); + _ = VectorSearcher.SelectTopKFromScores(scores, topK, float.NegativeInfinity); heapStopwatch.Stop(); heapMs = Math.Min(heapMs, heapStopwatch.Elapsed.TotalMilliseconds); Stopwatch fullSortStopwatch = Stopwatch.StartNew(); - reference = FullSortReferenceSearch(vectors, dimensions, query, topK); + _ = FullSortSelectTopK(scores, topK); fullSortStopwatch.Stop(); fullSortMs = Math.Min(fullSortMs, fullSortStopwatch.Elapsed.TotalMilliseconds); } _output.WriteLine( - $"VectorSearcher.Search (bounded heap) over {count} x {dimensions} vectors: " + - $"best of {measuredRuns} runs took {heapMs:F3} ms."); + $"VectorSearcher.SelectTopKFromScores (bounded heap) over {count} precomputed scores: " + + $"best of {measuredRuns} runs took {heapMs:F4} ms."); _output.WriteLine( - $"Reference full-sort search over the same data: " + - $"best of {measuredRuns} runs took {fullSortMs:F3} ms."); - - Assert.Equal(topK, hits.Count); - Assert.Equal(reference, hits); - - // The bound is relative, not absolute. An absolute millisecond threshold cannot work - // on a shared CI runner: the same search measured 1.6 ms locally and 3 s on a hosted - // Windows agent, a 2000x spread caused by virtualisation and noisy neighbours, not by - // any regression. Comparing against a full sort of the same data on the same machine - // in the same run cancels that out — both paths pay the identical environment tax. + $"Full sort of the same precomputed scores: " + + $"best of {measuredRuns} runs took {fullSortMs:F4} ms."); + + // The bound is relative, not absolute, for the same reason as before: an absolute + // millisecond threshold cannot work on a shared CI runner. Comparing the two selection + // strategies against the same fixed scores array in the same run cancels out the + // environment tax that made an absolute threshold unworkable, without also cancelling + // out (and thereby hiding, per issue #32) a genuine regression in the selection code — + // there is no shared dot-product cost left in this measurement to hide behind. // - // The claim under test is that the bounded heap is not slower than sorting everything, - // which is the entire justification for its extra complexity. Locally that ratio is - // about 0.65; the 1.5x ceiling tolerates measurement noise while still failing loudly - // if selection ever degrades into something worse than the naive approach. + // With the shared cost gone, local measurement clusters tightly around 0.02-0.03 (the + // heap does O(N log 20) work against the full sort's O(N log N), so it should win by a + // wide margin, not a narrow one) — 35 consecutive runs here, solo and under the full + // suite's default parallelism, never exceeded 0.03. A ceiling of 1.0 keeps the assertion + // legible as exactly the claim under test ("the heap must not be slower than sorting + // everything") while leaving roughly 30x headroom over the observed value for a slower or + // noisier CI runner — tight enough to fail loudly on a real regression, unlike the 1.5 + // ceiling this replaces, which tolerated noise from cost this measurement no longer pays. double ratio = heapMs / fullSortMs; - _output.WriteLine($"Heap/full-sort ratio: {ratio:F3} (lower is better)."); + _output.WriteLine($"Heap/full-sort selection ratio: {ratio:F3} (lower is better)."); Assert.True( - ratio < 1.5, - $"Bounded-heap selection took {heapMs:F3} ms against {fullSortMs:F3} ms for a full sort " + - $"(ratio {ratio:F3}). The heap is meant to be no slower than sorting everything."); + ratio < 1.0, + $"Bounded-heap selection took {heapMs:F4} ms against {fullSortMs:F4} ms for a full sort " + + $"of the same precomputed scores (ratio {ratio:F3}). The heap is meant to be no slower " + + $"than sorting everything."); } - private static ScoredIndex[] FullSortReferenceSearch(float[] vectors, int dimensions, float[] query, int topK) + /// + /// The naive baseline is compared against: + /// sort every score, then take the first . Operates on an already + /// -computed array (row index doubles as ) so the comparison is selection strategy against selection + /// strategy, not selection strategy against selection-strategy-plus-a-second-dot-product-pass. + /// + private static ScoredIndex[] FullSortSelectTopK(float[] scores, int topK) { - int count = vectors.Length / dimensions; + int count = scores.Length; ScoredIndex[] all = new ScoredIndex[count]; for (int i = 0; i < count; i++) - all[i] = new ScoredIndex(i, TensorPrimitives.Dot(vectors.AsSpan(i * dimensions, dimensions), query)); + all[i] = new ScoredIndex(i, scores[i]); Array.Sort(all, static (a, b) => {