From 46fc7f2c4ad61174a6b1cb2ad14889414b98caad Mon Sep 17 00:00:00 2001 From: David Snelling Date: Wed, 27 May 2026 13:13:47 -0700 Subject: [PATCH] feat: hook native sort:topK provider into search result ranking Adds a swappable sort:topK seam for the find() result-ranking hot path. Brainy ranks candidate results by relevance score then slices to offset+limit; for large candidate sets that is an O(N log N) full sort even though only the page is kept. New utils/resultRanking.ts exposes rankIndicesByScore (top-K index selection) with a JS default (sortTopKIndicesJs) that is byte-identical in ordering to the previous results.sort((a, b) => b.score - a.score) + slice: descending by score, stable ties by original index (ES2019+ stable sort semantics). setSortTopKImplementation lets a native provider (e.g. cortex's Rust partial-sort) replace it; the provider returns indices into a scores array and only has to match the JS index ordering. brainy.ts resolves the sort:topK provider in init() (same pattern as distance:sq8) and wires both relevance-score sort sites in find() through rankIndicesByScore + reorderByIndices. rankIndicesByScore validates a native provider's output (length, range, no duplicates) and falls back to JS on any inconsistency, so a query never returns wrong or duplicated results. No provider registered = unchanged JS behavior. Tests: unit coverage of the dispatcher (JS ordering vs Array.sort across 200 random trials incl. ties, provider routing, and fallback on bad/throwing providers) plus find() integration proving ranking routes through a registered provider, that the provider and JS fallback produce identical ids/scores on the same data, and that pagination requests k = offset + limit with descending ordering. --- src/brainy.ts | 32 +++- src/utils/resultRanking.ts | 160 +++++++++++++++++ tests/unit/utils/resultRanking.test.ts | 238 +++++++++++++++++++++++++ 3 files changed, 425 insertions(+), 5 deletions(-) create mode 100644 src/utils/resultRanking.ts create mode 100644 tests/unit/utils/resultRanking.test.ts diff --git a/src/brainy.ts b/src/brainy.ts index 368f6944..8819824e 100644 --- a/src/brainy.ts +++ b/src/brainy.ts @@ -35,6 +35,7 @@ import { createPipeline } from './streaming/pipeline.js' import { configureLogger, LogLevel } from './utils/logger.js' import { setGlobalCache } from './utils/unifiedCache.js' import type { UnifiedCache } from './utils/unifiedCache.js' +import { rankIndicesByScore, reorderByIndices } from './utils/resultRanking.js' import { PluginRegistry } from './plugin.js' import type { BrainyPlugin, BrainyPluginContext } from './plugin.js' import { TransactionManager } from './transaction/TransactionManager.js' @@ -510,6 +511,19 @@ export class Brainy implements BrainyInterface { setSQ8DistanceImplementation(sq8DistanceProvider) } + // Provider: sort:topK (e.g. cortex's native partial-sort / heap-select) — swaps the + // JS result-ranking used by find() to pick the top `offset + limit` rows. The provider + // returns indices into a scores array, ordered descending with stable ties, identical + // to the JS sortTopKIndicesJs. rankIndicesByScore validates the provider's output and + // falls back to JS on any inconsistency, so ranking is always correct. + const sortTopKProvider = this.pluginRegistry.getProvider< + (scores: number[], k: number, descending: boolean) => number[] + >('sort:topK') + if (sortTopKProvider) { + const { setSortTopKImplementation } = await import('./utils/resultRanking.js') + setSortTopKImplementation(sortTopKProvider) + } + // Provider: distance function (resolve BEFORE setupIndex — index uses this.distance) const nativeDistance = this.pluginRegistry.getProvider('distance') if (nativeDistance) { @@ -2585,10 +2599,14 @@ export class Brainy implements BrainyInterface { const limit = params.limit || 10 const offset = params.offset || 0 - // If we have enough filtered results, sort and paginate early + // If we have enough filtered results, sort and paginate early. + // Rank by score (top offset+limit), then drop the offset — identical ordering + // to a full `sort((a, b) => b.score - a.score)` + slice, but the native + // `sort:topK` provider can compute only the page instead of the full sort. if (results.length >= offset + limit) { - results.sort((a, b) => b.score - a.score) - results = results.slice(offset, offset + limit) + const k = offset + limit + const order = rankIndicesByScore(results.map(r => r.score), k, true) + results = reorderByIndices(results, order).slice(offset, k) // Batch-load entities only for the paginated results (10x faster on GCS) const idsToLoad = results.filter(r => !r.entity).map(r => r.id) @@ -2693,8 +2711,12 @@ export class Brainy implements BrainyInterface { results = resultsWithValues.map(({ result }) => result) } else { - // Default: sort by relevance score - results.sort((a, b) => b.score - a.score) + // Default: sort by relevance score. Rank to the page (offset+limit) via the + // swappable sort:topK seam; the slice below drops the offset. Ordering is + // identical to a full `sort((a, b) => b.score - a.score)`. + const k = (params.offset || 0) + limit + const order = rankIndicesByScore(results.map(r => r.score), k, true) + results = reorderByIndices(results, order) } const finalOffset = params.offset || 0 diff --git a/src/utils/resultRanking.ts b/src/utils/resultRanking.ts new file mode 100644 index 00000000..4186241e --- /dev/null +++ b/src/utils/resultRanking.ts @@ -0,0 +1,160 @@ +/** + * @module utils/resultRanking + * @description Top-K result ranking for the search pipeline. Brainy's `find()` ranks + * candidate results by relevance score and then slices to a page (`offset + limit`). + * For large candidate sets this is an O(N log N) full sort even though only the top + * `k = offset + limit` rows are kept. + * + * This module exposes a swappable `sort:topK` seam: + * - {@link rankIndicesByScore} returns the indices of `scores` ordered exactly as a + * descending stable sort would order them, then truncated to `k`. + * - {@link setSortTopKImplementation} lets a native plugin (e.g. cortex's Rust + * partial-sort / heap-select) replace the JS implementation. The native provider + * only has to return the same *index ordering* the JS path produces. + * + * The default JS implementation ({@link sortTopKIndicesJs}) is the canonical, always-correct + * fallback. It is intentionally identical in ordering to `Array.prototype.sort((a, b) => b - a)`: + * descending by score, with stable ordering for ties (lower original index first). This + * guarantees that wiring the hook in or out never changes which results a query returns or + * in what order — only how fast the ranking is computed. + * + * @see setSQ8DistanceImplementation in ./vectorQuantization.js for the sibling provider-swap pattern. + */ + +/** + * Signature of a `sort:topK` implementation. Given a parallel array of `scores`, a + * target count `k`, and a sort direction, it returns the indices into `scores` in the + * desired order, truncated to at most `k` entries. + * + * Contract (the native implementation MUST match the JS default exactly): + * - `descending: true` → indices ordered by score high→low. + * - Ties (equal scores) keep their original relative order (stable). + * - The returned array has `min(k, scores.length)` entries; `k <= 0` returns `[]`. + * + * @param scores - Relevance scores, one per candidate result. + * @param k - Maximum number of indices to return (the page size, `offset + limit`). + * @param descending - When true, highest score first (the only mode used by ranking). + * @returns Indices into `scores`, ordered and truncated to `k`. + */ +export type SortTopKFn = (scores: number[], k: number, descending: boolean) => number[] + +/** + * Pure-JS top-K index selection. Produces the same ordering as + * `scores.map((_, i) => i).sort((a, b) => descending ? scores[b] - scores[a] : scores[a] - scores[b])` + * and then truncates to `k`, but without allocating intermediate result objects. + * + * Stability: ties are broken by original index (ascending), matching the ES2019+ stable + * `Array.prototype.sort`. This is what keeps the hook's output byte-identical to the + * previous `results.sort((a, b) => b.score - a.score)` followed by `slice`. + * + * @param scores - Relevance scores, one per candidate result. + * @param k - Maximum number of indices to return. + * @param descending - When true, highest score first. + * @returns Indices into `scores`, ordered and truncated to `k`. + * @example + * sortTopKIndicesJs([0.1, 0.9, 0.5], 2, true) // [1, 2] (0.9 then 0.5) + */ +export function sortTopKIndicesJs(scores: number[], k: number, descending: boolean): number[] { + const n = scores.length + if (n === 0 || k <= 0) return [] + + // Build the index list and sort it the same way the consumer's comparator did. + // Sorting indices (not value/object pairs) keeps allocation minimal while preserving + // the exact comparator semantics, including stable tie-breaking by original index. + const indices = new Array(n) + for (let i = 0; i < n; i++) indices[i] = i + + indices.sort((a, b) => { + const sa = scores[a] + const sb = scores[b] + if (sa === sb) return a - b // stable: original order for ties + return descending ? sb - sa : sa - sb + }) + + return k >= n ? indices : indices.slice(0, k) +} + +/** + * Active `sort:topK` implementation. Defaults to {@link sortTopKIndicesJs}; replaced by a + * native implementation via {@link setSortTopKImplementation} when a `sort:topK` provider + * is registered. The native implementation must return the identical index ordering. + */ +let activeSortTopK: SortTopKFn = sortTopKIndicesJs + +/** + * Rank candidate indices by relevance score, dispatched to the active `sort:topK` + * implementation (JS by default, native when a provider is registered). + * + * This is the single entry point the search pipeline calls. It is defensive about a + * misbehaving native provider: the returned indices are validated to be the right length, + * in range, and free of duplicates; if the provider returns anything inconsistent the JS + * implementation is used instead so a query never returns wrong or duplicated results. + * + * @param scores - Relevance scores, one per candidate result. + * @param k - Maximum number of indices to return (the page size, `offset + limit`). + * @param descending - When true, highest score first (always true for relevance ranking). + * @returns Indices into `scores`, ordered and truncated to at most `k`. + */ +export function rankIndicesByScore(scores: number[], k: number, descending: boolean): number[] { + const n = scores.length + if (n === 0 || k <= 0) return [] + + if (activeSortTopK === sortTopKIndicesJs) { + return sortTopKIndicesJs(scores, k, descending) + } + + // A native provider is installed. Run it, then validate its output. The validation + // is O(result) — cheap relative to the sort it replaces — and guarantees correctness + // even if a provider has a bug: any inconsistency falls back to the trusted JS path. + const expectedLen = Math.min(k, n) + let candidate: number[] + try { + candidate = activeSortTopK(scores, k, descending) + } catch { + return sortTopKIndicesJs(scores, k, descending) + } + + if (!Array.isArray(candidate) || candidate.length !== expectedLen) { + return sortTopKIndicesJs(scores, k, descending) + } + + const seen = new Uint8Array(n) + for (let i = 0; i < candidate.length; i++) { + const idx = candidate[i] + if (!Number.isInteger(idx) || idx < 0 || idx >= n || seen[idx] === 1) { + return sortTopKIndicesJs(scores, k, descending) + } + seen[idx] = 1 + } + + return candidate +} + +/** + * Reorder `items` according to a top-K index ranking, returning a new array containing + * only the ranked (and truncated) elements. Used by the search pipeline to apply a + * {@link rankIndicesByScore} result to the parallel array of full result objects. + * + * @typeParam T - The result element type. + * @param items - The full candidate array (parallel to the `scores` passed to ranking). + * @param order - Indices into `items`, as returned by {@link rankIndicesByScore}. + * @returns A new array `[items[order[0]], items[order[1]], ...]`. + */ +export function reorderByIndices(items: T[], order: number[]): T[] { + const out = new Array(order.length) + for (let i = 0; i < order.length; i++) { + out[i] = items[order[i]] + } + return out +} + +/** + * Replace the `sort:topK` implementation at runtime (e.g. cortex's native partial sort). + * Called by `brainy.ts` when a `sort:topK` provider is registered. Pass + * {@link sortTopKIndicesJs} to restore the JS default. + * + * @param fn - Native implementation; must satisfy the {@link SortTopKFn} contract. + */ +export function setSortTopKImplementation(fn: SortTopKFn): void { + activeSortTopK = fn +} diff --git a/tests/unit/utils/resultRanking.test.ts b/tests/unit/utils/resultRanking.test.ts new file mode 100644 index 00000000..a2b8988b --- /dev/null +++ b/tests/unit/utils/resultRanking.test.ts @@ -0,0 +1,238 @@ +/** + * @module tests/unit/utils/resultRanking + * @description Tests for the `sort:topK` result-ranking hook. + * + * Two layers: + * 1. Unit — the swappable dispatcher (`rankIndicesByScore` / `setSortTopKImplementation`): + * JS default ordering matches `Array.prototype.sort`, a native provider is routed + * through, and a misbehaving provider falls back to JS. + * 2. Integration — `brain.find()` ranking routes through a registered `sort:topK` + * provider, AND the JS fallback (no provider) produces identical results/order. + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest' +import { + rankIndicesByScore, + reorderByIndices, + sortTopKIndicesJs, + setSortTopKImplementation, + type SortTopKFn +} from '../../../src/utils/resultRanking.js' +import { Brainy, NounType } from '../../../src/index.js' +import type { BrainyPlugin } from '../../../src/plugin.js' + +/** + * Reference ranking: exactly what the old `results.sort((a, b) => b.score - a.score)` + * followed by `slice(0, k)` produced. Used to assert the hook never changes ordering. + */ +function referenceTopK(scores: number[], k: number, descending: boolean): number[] { + const idx = scores.map((_, i) => i) + idx.sort((a, b) => (descending ? scores[b] - scores[a] : scores[a] - scores[b]) || a - b) + return idx.slice(0, Math.min(k, scores.length)) +} + +describe('resultRanking — sortTopKIndicesJs (canonical JS ranking)', () => { + it('orders descending by score and truncates to k', () => { + const scores = [0.1, 0.9, 0.5, 0.7, 0.3] + expect(sortTopKIndicesJs(scores, 3, true)).toEqual([1, 3, 2]) // 0.9, 0.7, 0.5 + }) + + it('orders ascending when descending=false', () => { + const scores = [0.1, 0.9, 0.5] + expect(sortTopKIndicesJs(scores, 3, false)).toEqual([0, 2, 1]) // 0.1, 0.5, 0.9 + }) + + it('breaks ties by original index (stable), matching Array.sort', () => { + // Three equal scores plus a higher one. Stable order keeps 0,2,3 in input order. + const scores = [0.5, 0.9, 0.5, 0.5] + expect(sortTopKIndicesJs(scores, 4, true)).toEqual([1, 0, 2, 3]) + }) + + it('returns [] for empty input or non-positive k', () => { + expect(sortTopKIndicesJs([], 5, true)).toEqual([]) + expect(sortTopKIndicesJs([0.1, 0.2], 0, true)).toEqual([]) + expect(sortTopKIndicesJs([0.1, 0.2], -1, true)).toEqual([]) + }) + + it('returns all indices when k exceeds length', () => { + expect(sortTopKIndicesJs([0.2, 0.1], 100, true)).toEqual([0, 1]) + }) + + it('matches the reference Array.sort ordering across random inputs', () => { + for (let trial = 0; trial < 200; trial++) { + const n = 1 + Math.floor(Math.random() * 40) + // Include duplicates so tie-stability is exercised. + const scores = Array.from({ length: n }, () => Math.floor(Math.random() * 5) / 5) + const k = 1 + Math.floor(Math.random() * (n + 2)) + expect(sortTopKIndicesJs(scores, k, true)).toEqual(referenceTopK(scores, k, true)) + } + }) +}) + +describe('resultRanking — reorderByIndices', () => { + it('reorders elements according to the index order', () => { + const items = ['a', 'b', 'c', 'd'] + expect(reorderByIndices(items, [2, 0, 3])).toEqual(['c', 'a', 'd']) + }) + + it('returns an empty array for an empty order', () => { + expect(reorderByIndices(['a', 'b'], [])).toEqual([]) + }) +}) + +describe('resultRanking — rankIndicesByScore dispatcher + provider hook', () => { + // Always restore the JS default so a swapped impl never leaks into other tests. + afterEach(() => setSortTopKImplementation(sortTopKIndicesJs)) + + it('uses the JS implementation by default (identical to reference)', () => { + const scores = [0.4, 0.8, 0.1, 0.8] + expect(rankIndicesByScore(scores, 3, true)).toEqual(referenceTopK(scores, 3, true)) + }) + + it('routes through a registered sort:topK provider with the documented signature', () => { + const calls: Array<[number[], number, boolean]> = [] + const provider: SortTopKFn = (scores, k, descending) => { + calls.push([scores, k, descending]) + // Return a valid (correct) ranking so the dispatcher accepts it. + return sortTopKIndicesJs(scores, k, descending) + } + setSortTopKImplementation(provider) + + const scores = [0.2, 0.9, 0.5] + const order = rankIndicesByScore(scores, 2, true) + + expect(calls).toHaveLength(1) + expect(calls[0][0]).toBe(scores) // exact scores array + expect(calls[0][1]).toBe(2) // k + expect(calls[0][2]).toBe(true) // descending + expect(order).toEqual([1, 2]) // 0.9, 0.5 + }) + + it('falls back to JS when a provider returns the wrong length', () => { + setSortTopKImplementation(() => [0]) // wrong length (expected 2) + const scores = [0.2, 0.9, 0.5] + expect(rankIndicesByScore(scores, 2, true)).toEqual(referenceTopK(scores, 2, true)) + }) + + it('falls back to JS when a provider returns out-of-range or duplicate indices', () => { + setSortTopKImplementation(() => [0, 99]) // 99 out of range + expect(rankIndicesByScore([0.2, 0.9, 0.5], 2, true)).toEqual([1, 2]) + + setSortTopKImplementation(() => [1, 1]) // duplicate index + expect(rankIndicesByScore([0.2, 0.9, 0.5], 2, true)).toEqual([1, 2]) + }) + + it('falls back to JS when a provider throws', () => { + setSortTopKImplementation(() => { + throw new Error('native boom') + }) + expect(rankIndicesByScore([0.2, 0.9, 0.5], 2, true)).toEqual([1, 2]) + }) +}) + +/** + * Plugin that registers a recording `sort:topK` provider. The provider delegates to the + * correct JS ranking so results are valid, while recording that it was invoked — proving + * find()'s ranking routes through the hook. + */ +function makeSortTopKPlugin(calls: Array<{ k: number; descending: boolean; n: number }>): BrainyPlugin { + return { + name: 'test-sort-topk', + activate: async (ctx) => { + ctx.registerProvider('sort:topK', (scores: number[], k: number, descending: boolean) => { + calls.push({ k, descending, n: scores.length }) + return sortTopKIndicesJs(scores, k, descending) + }) + return true + } + } +} + +describe('resultRanking — Brainy.find() integration', () => { + // Restore JS default after every test: brain.init() installs the provider process-wide + // via setSortTopKImplementation, so we must reset it to avoid leaking into other suites. + afterEach(() => setSortTopKImplementation(sortTopKIndicesJs)) + + /** Seed a brain with deterministic data covering several types. */ + async function seed(brain: Brainy): Promise { + const docs = [ + { data: 'red apple fruit', type: NounType.Concept, metadata: { name: 'apple' } }, + { data: 'green pear fruit', type: NounType.Concept, metadata: { name: 'pear' } }, + { data: 'yellow banana fruit', type: NounType.Concept, metadata: { name: 'banana' } }, + { data: 'orange citrus fruit', type: NounType.Concept, metadata: { name: 'orange' } }, + { data: 'purple grape fruit', type: NounType.Concept, metadata: { name: 'grape' } }, + { data: 'blue car vehicle', type: NounType.Thing, metadata: { name: 'car' } }, + { data: 'fast train vehicle', type: NounType.Thing, metadata: { name: 'train' } } + ] + for (const d of docs) await brain.add(d) + } + + it('routes find() relevance ranking through a registered sort:topK provider', async () => { + const calls: Array<{ k: number; descending: boolean; n: number }> = [] + const brain = new Brainy({ storage: { type: 'memory' }, silent: true }) + brain.use(makeSortTopKPlugin(calls)) + await brain.init() + await seed(brain) + + const results = await brain.find({ query: 'fruit', limit: 3 }) + + // The provider must have been invoked for relevance ranking, with descending order + // and k == offset(0) + limit(3). + expect(calls.length).toBeGreaterThan(0) + const ranking = calls.find(c => c.descending && c.k === 3) + expect(ranking).toBeDefined() + expect(results.length).toBeGreaterThan(0) + expect(results.length).toBeLessThanOrEqual(3) + + await brain.close() + }) + + it('produces identical results via the provider and the JS fallback (same brain, same data)', async () => { + // The active sort:topK implementation is a process-global, so the only sound way to + // compare provider-vs-JS is on ONE brain over identical persisted data: run with the + // provider active, then reset to JS and re-run the same query. The candidate set and + // scores are identical across runs — only the ranking code path differs — so ids, + // order, and scores must match exactly. (Two separate brains would differ due to + // independent approximate-NN candidate selection, not the ranking hook.) + const calls: Array<{ k: number; descending: boolean; n: number }> = [] + const brain = new Brainy({ storage: { type: 'memory' }, silent: true }) + brain.use(makeSortTopKPlugin(calls)) + await brain.init() + await seed(brain) + + const query = { query: 'fruit', limit: 4 } + const providerResults = await brain.find({ ...query }) + expect(calls.length).toBeGreaterThan(0) + + // Now disable the provider (restore JS) and re-run the identical query on the same brain. + setSortTopKImplementation(sortTopKIndicesJs) + const jsResults = await brain.find({ ...query }) + + expect(providerResults.map(r => r.id)).toEqual(jsResults.map(r => r.id)) + expect(providerResults.map(r => r.score)).toEqual(jsResults.map(r => r.score)) + + await brain.close() + }) + + it('returns results in non-increasing score order and requests k = offset + limit', async () => { + // Ordering correctness is the hook's contract; assert it directly on find() output + // rather than relying on cross-query ANN stability. Also assert the provider is asked + // for the full page (offset + limit), which is what makes pagination correct. + const calls: Array<{ k: number; descending: boolean; n: number }> = [] + const brain = new Brainy({ storage: { type: 'memory' }, silent: true }) + brain.use(makeSortTopKPlugin(calls)) + await brain.init() + await seed(brain) + + const page = await brain.find({ query: 'fruit', limit: 2, offset: 2 }) + + // Scores within the returned page must be non-increasing (descending ranking). + for (let i = 1; i < page.length; i++) { + expect(page[i].score).toBeLessThanOrEqual(page[i - 1].score) + } + // The offset page must request k = offset(2) + limit(2) = 4 from the provider. + expect(calls.some(c => c.descending && c.k === 4)).toBe(true) + + await brain.close() + }) +})