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.
This commit is contained in:
parent
00d14cfa93
commit
46fc7f2c4a
3 changed files with 425 additions and 5 deletions
160
src/utils/resultRanking.ts
Normal file
160
src/utils/resultRanking.ts
Normal file
|
|
@ -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<number>(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<T>(items: T[], order: number[]): T[] {
|
||||
const out = new Array<T>(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
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue