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:
David Snelling 2026-05-27 13:13:47 -07:00
parent 00d14cfa93
commit 46fc7f2c4a
3 changed files with 425 additions and 5 deletions

View file

@ -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<T = any> implements BrainyInterface<T> {
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<DistanceFunction>('distance')
if (nativeDistance) {
@ -2585,10 +2599,14 @@ export class Brainy<T = any> implements BrainyInterface<T> {
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<T = any> implements BrainyInterface<T> {
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