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
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue