feat: mmap-vector backend wiring — HNSWIndex consumes vectorStore:mmap (2.4.0 #2)
Cortex already registers the vectorStore:mmap provider (its Rust NativeMmapVectorStore), but brainy has never consumed it — preloadVectors and getVectorSafe still go straight to storage.getNounVector for every id, even when an mmap layer is available. This wires the consumer end. Architecture: - NEW MmapVectorBackend (src/hnsw/mmapVectorBackend.ts) — bridges brainy's UUID-keyed vector reads to a int-slot mmap file via the vectorStore:mmap provider. Slots are addressed by the stable int id from the post-2.4.0 #1 EntityIdMapper (the foundation this depends on). Auto-grows the file (doubling) when a write lands beyond capacity, so HNSWIndex never has to think about sizing. The class never touches per-entity storage — it owns only the mmap layer. - HNSWIndex changes — adds a vectorBackend field + a setVectorBackend setter. The vector read paths (preloadVectors, getVectorSafe) try the mmap layer first; on a storage fallback hit, they LAZILY write back into the mmap slot. An upgraded install converges to the zero-copy fast path under live traffic — no big-bang migration step. The legacy per-entity path is preserved and still used when no backend is set. - brainy.ts wiring — a new private wireMmapVectorBackend() runs once during init, after plugin activation + metadataIndex setup. It activates the backend only when (a) the vectorStore:mmap provider is registered, (b) the storage adapter resolves a real local path via getBinaryBlobPath(), and (c) the metadata index exposes its idMapper. Cloud adapters return null on (b) and the backend is silently skipped; HNSWIndex's behaviour is then identical to pre-2.4.0. - Provider interfaces in plugin.ts — VectorStoreMmapProvider and VectorStoreMmapInstance document the contract cortex's class fulfils (the class IS the provider — static factory methods). Brainy depends on the interfaces, not on cortex; the structural match is verified when cortex 2.4.0 picks up this brainy release. Tests (1428 total, +11 vs pre-2.4.0): - tests/unit/hnsw/mmap-vector-backend.test.ts — 6 unit tests with an in-memory mock provider. Covers round-trip, batch reads with interleaved misses, slot stability (no re-slotting on overwrite), file growth without data loss, idempotent open, and null returns for unwritten slots. The real perf integration with cortex's NativeMmapVectorStore is exercised when cortex 2.4.0 wires this in. - tests/unit/utils/entity-id-mapper-stability.test.ts — moved here from tests/regression/ (which is NOT in the unit-config include glob, so the five #23 tests were not actually being run by npm test). The unit config matches tests/unit/**/*.test.ts. The 2.4.0 #2 follow-up will be the chunked-segment layout for remote storage adapters (S3 / R2 / GCS) where a single growing file doesn't fit immutable objects. For 2.4.0 release: local-FS only.
This commit is contained in:
parent
b2408cb5a6
commit
d4cb26c604
6 changed files with 640 additions and 4 deletions
|
|
@ -37,7 +37,8 @@ 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 type { BrainyPlugin, BrainyPluginContext, VectorStoreMmapProvider } from './plugin.js'
|
||||
import { MmapVectorBackend } from './hnsw/mmapVectorBackend.js'
|
||||
import { TransactionManager } from './transaction/TransactionManager.js'
|
||||
import {
|
||||
ValidationConfig,
|
||||
|
|
@ -559,6 +560,15 @@ export class Brainy<T = any> implements BrainyInterface<T> {
|
|||
this.graphIndex = graphIndex
|
||||
}
|
||||
|
||||
// Wire the mmap-vector backend (2.4.0 #2). When the vectorStore:mmap
|
||||
// provider is registered AND the storage adapter resolves a real local
|
||||
// path via getBinaryBlobPath(), open the mmap file there and inject the
|
||||
// backend into HNSWIndex. Cloud adapters return null and skip silently;
|
||||
// HNSWIndex's behaviour is then identical to pre-2.4.0 (per-entity reads
|
||||
// from canonical storage). Failures are non-fatal — the index still
|
||||
// works, just without the zero-copy fast path.
|
||||
await this.wireMmapVectorBackend()
|
||||
|
||||
// Rebuild indexes if needed for existing data
|
||||
await this.rebuildIndexesIfNeeded()
|
||||
|
||||
|
|
@ -7716,6 +7726,72 @@ export class Brainy<T = any> implements BrainyInterface<T> {
|
|||
await this.lazyRebuildPromise
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Open the mmap-vector backend and inject it into HNSWIndex.
|
||||
* Called once during init after plugin activation, metadataIndex init, and
|
||||
* graphIndex setup — so the idMapper is available and providers are wired.
|
||||
*
|
||||
* Activation conditions (ALL must hold; otherwise this is a no-op and the
|
||||
* index falls back to per-entity storage reads, the pre-2.4.0 behaviour):
|
||||
* 1. The `vectorStore:mmap` provider is registered (cortex registers this).
|
||||
* 2. The storage adapter resolves a real local path via `getBinaryBlobPath()`
|
||||
* — cloud adapters return null and the mmap layer is skipped silently.
|
||||
* 3. The metadata index exposes its idMapper (a stable UUID↔int map). The
|
||||
* mmap layout is keyed by these ints, so 2.4.0 #1 is the prerequisite.
|
||||
*
|
||||
* Vector dimensionality: read from `this.dimensions` if set; otherwise the
|
||||
* default for the standard embedding model (384). The dim is fixed at
|
||||
* file-creation time and cannot change later, so a wrong default would only
|
||||
* matter on the very first init of a brand-new instance. After the first
|
||||
* `add()` the in-memory `this.dimensions` is set, and any subsequent init
|
||||
* (or a re-open) sees the correct value.
|
||||
*
|
||||
* Failures here are non-fatal: a debug-level log records the reason, and
|
||||
* HNSWIndex continues without the mmap fast path.
|
||||
*/
|
||||
private async wireMmapVectorBackend(): Promise<void> {
|
||||
const provider = this.pluginRegistry.getProvider<VectorStoreMmapProvider>('vectorStore:mmap')
|
||||
if (!provider) return
|
||||
|
||||
const storageWithBlob = this.storage as unknown as {
|
||||
getBinaryBlobPath?: (key: string) => string | null
|
||||
}
|
||||
const vectorPath = storageWithBlob.getBinaryBlobPath?.('_vectors/main') ?? null
|
||||
if (!vectorPath) return
|
||||
|
||||
const idMapper = this.metadataIndex.getIdMapper?.()
|
||||
if (!idMapper) return
|
||||
|
||||
const dim = this.dimensions ?? 384
|
||||
const initialCapacity = Math.max(idMapper.size * 2, 1024)
|
||||
|
||||
try {
|
||||
const backend = await MmapVectorBackend.open(
|
||||
provider,
|
||||
vectorPath,
|
||||
dim,
|
||||
initialCapacity,
|
||||
idMapper
|
||||
)
|
||||
this.index.setVectorBackend(backend)
|
||||
if (!this.config.silent) {
|
||||
console.log(
|
||||
`[brainy] vector mmap backend wired (path=${vectorPath}, dim=${dim}, capacity=${initialCapacity})`
|
||||
)
|
||||
}
|
||||
} catch (error) {
|
||||
// Common cases: dim mismatch from a prior init at a different dimension,
|
||||
// permission errors, disk full, file corruption. Index keeps working via
|
||||
// the per-entity storage path; flag the cause for diagnosis.
|
||||
if (!this.config.silent) {
|
||||
console.warn(
|
||||
`[brainy] mmap-vector backend not wired (${(error as Error).message}); ` +
|
||||
`falling back to per-entity vector reads`
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Rebuild indexes from persisted data if needed (LAZY LOADING)
|
||||
*
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import { prodLog } from '../utils/logger.js'
|
|||
import { quantizeSQ8, distanceSQ8 } from '../utils/vectorQuantization.js'
|
||||
import type { SQ8QuantizedVector } from '../utils/vectorQuantization.js'
|
||||
import type { HnswProvider } from '../plugin.js'
|
||||
import { MmapVectorBackend } from './mmapVectorBackend.js'
|
||||
|
||||
// Default HNSW parameters
|
||||
const DEFAULT_CONFIG: HNSWConfig = {
|
||||
|
|
@ -49,6 +50,13 @@ export class HNSWIndex implements HnswProvider {
|
|||
private unifiedCache: UnifiedCache // Shared cache with Graph and Metadata indexes
|
||||
// Always-adaptive caching - no "mode" concept, system adapts automatically
|
||||
|
||||
// Optional mmap-vector backend (2.4.0 #2). When set, the read paths
|
||||
// (preloadVectors, getVectorSafe) try the mmap layer before storage, and
|
||||
// on a storage hit they write back into the mmap slot — a lazy migration
|
||||
// that converges upgraded installs to the fast path under live traffic.
|
||||
// Null on cloud storage adapters (no local path) and pre-injection init.
|
||||
private vectorBackend: MmapVectorBackend | null = null
|
||||
|
||||
// COW (Copy-on-Write) support
|
||||
private cowEnabled: boolean = false
|
||||
private cowModifiedNodes: Set<string> = new Set()
|
||||
|
|
@ -101,6 +109,22 @@ export class HNSWIndex implements HnswProvider {
|
|||
this.useParallelization = useParallelization
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Inject (or detach) the mmap-vector backend. When set, the
|
||||
* vector read paths (`preloadVectors`, `getVectorSafe`) try the mmap layer
|
||||
* before storage and lazily write back on storage hits — converging an
|
||||
* upgraded install to the zero-copy fast path without a migration step.
|
||||
* Setting to null reverts to the existing per-entity read path.
|
||||
*
|
||||
* Lifecycle: brainy.ts wires this after plugin activation when (a) the
|
||||
* `vectorStore:mmap` provider is registered AND (b) the storage adapter
|
||||
* resolves a real local path via `getBinaryBlobPath()`. Cloud adapters
|
||||
* leave it null; HNSWIndex's behaviour is then identical to pre-2.4.0.
|
||||
*/
|
||||
public setVectorBackend(backend: MmapVectorBackend | null): void {
|
||||
this.vectorBackend = backend
|
||||
}
|
||||
|
||||
/**
|
||||
* Get whether parallelization is enabled
|
||||
*/
|
||||
|
|
@ -1073,7 +1097,23 @@ export class HNSWIndex implements HnswProvider {
|
|||
const cacheKey = `hnsw:vector:${noun.id}`
|
||||
|
||||
const vector = await this.unifiedCache.get(cacheKey, async () => {
|
||||
// Cache miss - load from storage
|
||||
// Try the mmap-vector backend first when wired — zero-copy slot read,
|
||||
// no per-entity object hydration.
|
||||
if (this.vectorBackend) {
|
||||
const fromMmap = this.vectorBackend.readByUuid(noun.id)
|
||||
if (fromMmap) {
|
||||
this.unifiedCache.set(
|
||||
cacheKey,
|
||||
fromMmap,
|
||||
'hnsw',
|
||||
fromMmap.length * 4,
|
||||
50
|
||||
)
|
||||
return fromMmap
|
||||
}
|
||||
}
|
||||
|
||||
// Storage fallback — the canonical source of truth.
|
||||
if (!this.storage) {
|
||||
throw new Error('Storage not available for vector loading')
|
||||
}
|
||||
|
|
@ -1083,6 +1123,18 @@ export class HNSWIndex implements HnswProvider {
|
|||
throw new Error(`Vector not found for noun ${noun.id}`)
|
||||
}
|
||||
|
||||
// Lazy migration: write back into the mmap slot so the next read is
|
||||
// mmap-fast. Idempotent + resumable; a write failure is non-fatal —
|
||||
// the field still serves this read, and the next storage hit re-tries
|
||||
// the migration. (See mmapVectorBackend.ts for the contract.)
|
||||
if (this.vectorBackend) {
|
||||
try {
|
||||
this.vectorBackend.writeByUuid(noun.id, loaded)
|
||||
} catch (error) {
|
||||
prodLog.debug(`MmapVectorBackend write-back failed for ${noun.id}`, error)
|
||||
}
|
||||
}
|
||||
|
||||
// Add to UnifiedCache with cost-aware eviction
|
||||
// This competes fairly with Graph and Metadata indexes
|
||||
this.unifiedCache.set(
|
||||
|
|
@ -1140,7 +1192,50 @@ export class HNSWIndex implements HnswProvider {
|
|||
private async preloadVectors(nodeIds: string[]): Promise<void> {
|
||||
if (nodeIds.length === 0) return
|
||||
|
||||
// Use UnifiedCache's request coalescing to prevent duplicate loads
|
||||
// mmap-vector backend fast path: batched O(1) slot reads + lazy
|
||||
// write-back migration on miss. Cache populated from both layers so
|
||||
// subsequent in-session reads are O(1) memory.
|
||||
if (this.vectorBackend) {
|
||||
// Filter out ids already cached so we never re-read.
|
||||
const uncachedIds: string[] = []
|
||||
for (const id of nodeIds) {
|
||||
if (this.unifiedCache.getSync(`hnsw:vector:${id}`) === undefined) {
|
||||
uncachedIds.push(id)
|
||||
}
|
||||
}
|
||||
if (uncachedIds.length === 0) return
|
||||
|
||||
const fromMmap = this.vectorBackend.readBatchByUuid(uncachedIds)
|
||||
const misses: string[] = []
|
||||
for (let i = 0; i < uncachedIds.length; i++) {
|
||||
const id = uncachedIds[i]
|
||||
const v = fromMmap[i]
|
||||
if (v) {
|
||||
this.unifiedCache.set(`hnsw:vector:${id}`, v, 'hnsw', v.length * 4, 50)
|
||||
} else {
|
||||
misses.push(id)
|
||||
}
|
||||
}
|
||||
if (misses.length === 0) return
|
||||
|
||||
// Storage fallback for the misses, with concurrent lazy write-back.
|
||||
await Promise.all(misses.map(async (id) => {
|
||||
if (!this.storage) return
|
||||
const vector = await this.storage.getNounVector(id)
|
||||
if (!vector) return
|
||||
if (this.vectorBackend) {
|
||||
try {
|
||||
this.vectorBackend.writeByUuid(id, vector)
|
||||
} catch (error) {
|
||||
prodLog.debug(`MmapVectorBackend write-back failed for ${id}`, error)
|
||||
}
|
||||
}
|
||||
this.unifiedCache.set(`hnsw:vector:${id}`, vector, 'hnsw', vector.length * 4, 50)
|
||||
}))
|
||||
return
|
||||
}
|
||||
|
||||
// Legacy path: per-entity storage load via UnifiedCache request coalescing.
|
||||
const promises = nodeIds.map(async (id) => {
|
||||
const cacheKey = `hnsw:vector:${id}`
|
||||
return this.unifiedCache.get(cacheKey, async () => {
|
||||
|
|
|
|||
174
src/hnsw/mmapVectorBackend.ts
Normal file
174
src/hnsw/mmapVectorBackend.ts
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
/**
|
||||
* @module hnsw/mmapVectorBackend
|
||||
* @description Disk-resident vector cache backing HNSWIndex, wrapping the
|
||||
* `vectorStore:mmap` provider (cortex's NativeMmapVectorStore in production).
|
||||
*
|
||||
* **Why this exists.** Brainy's per-entity vector storage hits two ceilings as
|
||||
* datasets grow: (1) every `HNSWIndex.preloadVectors` does a sequential
|
||||
* `storage.getNounVector` per id (no batch), and (2) cloud and FS storage both
|
||||
* serialize each vector through JSON / msgpack, blocking zero-copy reads. An
|
||||
* mmap-backed single-file vector store, addressed by the stable int from
|
||||
* `EntityIdMapper`, gives O(1) batch reads + OS-page-cache prefetch and is the
|
||||
* foundation 2.4.0 #2 needs.
|
||||
*
|
||||
* **Lazy migration on miss.** First read of an unmigrated UUID falls back to
|
||||
* `storage.getNounVector` and writes the bytes into the mmap slot before
|
||||
* returning — idempotent, resumable, and lets an upgraded install converge to
|
||||
* the fast path under live traffic without a big-bang migration step. The
|
||||
* per-entity store remains the canonical source of truth; nothing in the per-
|
||||
* entity layer is rewritten or deleted here.
|
||||
*
|
||||
* **Local-FS only in 2.4.0.** Activated when the storage adapter exposes a
|
||||
* real local path via `getBinaryBlobPath()`. Cloud adapters (S3 / R2 / GCS)
|
||||
* return null and the backend is not constructed, so HNSWIndex falls back to
|
||||
* the existing per-entity read path. Chunked vector segments for remote
|
||||
* storage are the planned 2.4.0 #2 follow-up.
|
||||
*/
|
||||
|
||||
import { mkdir } from 'node:fs/promises'
|
||||
import { dirname } from 'node:path'
|
||||
import type { Vector } from '../coreTypes.js'
|
||||
import type {
|
||||
EntityIdMapperProvider,
|
||||
VectorStoreMmapInstance,
|
||||
VectorStoreMmapProvider
|
||||
} from '../plugin.js'
|
||||
|
||||
/**
|
||||
* @description Bridge between brainy's UUID-keyed vector API and cortex's
|
||||
* int-slot mmap vector store. Translates UUIDs to stable int slots via the
|
||||
* `EntityIdMapper`, grows the file on demand, and handles read misses by
|
||||
* surfacing `null` (HNSWIndex's caller does the lazy storage fallback +
|
||||
* write-back). The class never touches per-entity storage directly — it owns
|
||||
* only the mmap layer.
|
||||
*/
|
||||
export class MmapVectorBackend {
|
||||
private readonly store: VectorStoreMmapInstance
|
||||
private readonly idMapper: EntityIdMapperProvider
|
||||
|
||||
private constructor(store: VectorStoreMmapInstance, idMapper: EntityIdMapperProvider) {
|
||||
this.store = store
|
||||
this.idMapper = idMapper
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Open or create the mmap vector file at `path`. If the file
|
||||
* does not exist, it is created with the given dim + initial capacity. If
|
||||
* the directory containing `path` does not exist, it is created recursively
|
||||
* (the cortex provider does not auto-create parent directories).
|
||||
* @param provider - The `vectorStore:mmap` provider (cortex
|
||||
* `NativeMmapVectorStore` in production).
|
||||
* @param path - Absolute path to the vector file. Derived by the caller
|
||||
* from `storage.getBinaryBlobPath()` on local-FS adapters; cloud adapters
|
||||
* return null and this method is not called.
|
||||
* @param dim - Vector dimensionality. Must match the embedding model.
|
||||
* @param initialCapacity - Slot capacity to allocate when creating a new
|
||||
* file. The file grows automatically on writes beyond capacity.
|
||||
* @param idMapper - Stable UUID↔int mapper (post-2.4.0 #1 semantics — never
|
||||
* renumbers across rebuild). The mmap slots are addressed by these ints,
|
||||
* so id stability is the load-bearing invariant of the layout.
|
||||
* @returns A ready-to-use backend.
|
||||
*/
|
||||
static async open(
|
||||
provider: VectorStoreMmapProvider,
|
||||
path: string,
|
||||
dim: number,
|
||||
initialCapacity: number,
|
||||
idMapper: EntityIdMapperProvider
|
||||
): Promise<MmapVectorBackend> {
|
||||
await mkdir(dirname(path), { recursive: true })
|
||||
let store: VectorStoreMmapInstance
|
||||
try {
|
||||
store = provider.open(path)
|
||||
} catch {
|
||||
store = provider.create(path, dim, Math.max(initialCapacity, 16))
|
||||
}
|
||||
return new MmapVectorBackend(store, idMapper)
|
||||
}
|
||||
|
||||
/** Vector dimensionality of the underlying mmap file. */
|
||||
get dim(): number {
|
||||
return this.store.dim
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Read a vector by UUID. Returns `null` for any of:
|
||||
* (1) UUID not yet assigned an int (entity never written through this
|
||||
* backend or storage), (2) int beyond the highest written slot, or
|
||||
* (3) the underlying mmap read throwing (corrupt / out-of-bounds).
|
||||
* The null path is the caller's signal to fall back to per-entity storage
|
||||
* and trigger a write-back migration.
|
||||
*/
|
||||
readByUuid(uuid: string): Vector | null {
|
||||
const intId = this.idMapper.getInt(uuid)
|
||||
if (intId === undefined || intId >= this.store.count) return null
|
||||
try {
|
||||
return this.store.readVector(intId) as unknown as Vector
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Batch read vectors by UUID. Returns an array aligned with
|
||||
* `uuids` where each entry is the vector or `null` for a miss. Order-
|
||||
* preserving. Internally collapses to a single batched mmap read of only the
|
||||
* indices that resolve to a written slot, then redistributes the flat output
|
||||
* back to the caller's order — so missing slots cost only a couple of map
|
||||
* lookups, not a full N-element read.
|
||||
*/
|
||||
readBatchByUuid(uuids: string[]): (Vector | null)[] {
|
||||
const indices: number[] = []
|
||||
const slotByUuid = new Map<string, number>()
|
||||
for (const uuid of uuids) {
|
||||
const intId = this.idMapper.getInt(uuid)
|
||||
if (intId !== undefined && intId < this.store.count) {
|
||||
slotByUuid.set(uuid, indices.length)
|
||||
indices.push(intId)
|
||||
}
|
||||
}
|
||||
if (indices.length === 0) return uuids.map(() => null)
|
||||
const flat = this.store.readVectorsBatch(indices)
|
||||
const dim = this.store.dim
|
||||
return uuids.map(uuid => {
|
||||
const slot = slotByUuid.get(uuid)
|
||||
if (slot === undefined) return null
|
||||
return flat.slice(slot * dim, (slot + 1) * dim) as unknown as Vector
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Write a vector at the slot for `uuid`. Assigns a new int if
|
||||
* the UUID is unseen (stable, append-only). Grows the underlying file (slot
|
||||
* capacity doubled until the target index fits) before writing past
|
||||
* capacity, so the caller need not worry about file sizing.
|
||||
*/
|
||||
writeByUuid(uuid: string, vector: Vector): void {
|
||||
const intId = this.idMapper.getOrAssign(uuid)
|
||||
if (intId >= this.store.capacity) {
|
||||
let newCap = Math.max(this.store.capacity, 16)
|
||||
while (intId >= newCap) newCap *= 2
|
||||
this.store.resize(newCap)
|
||||
}
|
||||
this.store.writeVector(intId, Array.from(vector as ArrayLike<number>))
|
||||
}
|
||||
|
||||
/**
|
||||
* @description madvise(WILLNEED) hint for the slots backing `uuids`. Cheap;
|
||||
* silently drops UUIDs not yet in the mmap. Used by HNSW search to overlap
|
||||
* disk I/O with graph traversal.
|
||||
*/
|
||||
prefetchByUuid(uuids: string[]): void {
|
||||
const indices: number[] = []
|
||||
for (const uuid of uuids) {
|
||||
const intId = this.idMapper.getInt(uuid)
|
||||
if (intId !== undefined && intId < this.store.count) indices.push(intId)
|
||||
}
|
||||
if (indices.length > 0) this.store.prefetch(indices)
|
||||
}
|
||||
|
||||
/** msync — flush dirty pages to disk. */
|
||||
flush(): void {
|
||||
this.store.flush()
|
||||
}
|
||||
}
|
||||
|
|
@ -253,6 +253,55 @@ export interface CacheProvider {
|
|||
// uses at the `getProvider('embeddings')` call site. No separate interface is
|
||||
// added here to avoid a duplicate, unwired contract.
|
||||
|
||||
/**
|
||||
* The `'vectorStore:mmap'` provider — a static factory class for an mmap-backed
|
||||
* vector file. Brainy's HNSWIndex uses the static `.create()` / `.open()`
|
||||
* factories to open a single-file store at a derived path on disk-resident
|
||||
* storage adapters; non-FS adapters with no local-path support skip the mmap
|
||||
* layer and fall back to per-entity vector reads.
|
||||
*
|
||||
* Cortex registers `NativeMmapVectorStore` as this provider — a Rust mmap'd
|
||||
* f32 array with a header, batch read/write APIs, and madvise prefetch. Slot
|
||||
* indices are the stable ints from `EntityIdMapper`; lazy migration on read
|
||||
* miss writes back from per-entity storage into the slot, converging an
|
||||
* upgraded install to the fast path without a big-bang step.
|
||||
*/
|
||||
export interface VectorStoreMmapProvider {
|
||||
/** Create a new vector file with pre-allocated slot capacity. */
|
||||
create(path: string, dim: number, capacity: number): VectorStoreMmapInstance
|
||||
/** Open an existing vector file for read + write. */
|
||||
open(path: string): VectorStoreMmapInstance
|
||||
/** Open an existing vector file for read-only access. */
|
||||
openReadOnly(path: string): VectorStoreMmapInstance
|
||||
}
|
||||
|
||||
/**
|
||||
* An open mmap-vector-store instance — the surface brainy's HNSWIndex consumes
|
||||
* for batch vector reads + lazy migration from per-entity storage on miss.
|
||||
*/
|
||||
export interface VectorStoreMmapInstance {
|
||||
/** Number of vectors written (highest written index + 1). */
|
||||
readonly count: number
|
||||
/** Maximum slot capacity before resize is needed. */
|
||||
readonly capacity: number
|
||||
/** Vector dimensionality. */
|
||||
readonly dim: number
|
||||
/** Write a single vector at `index`. Vector is f64 in JS, stored as f32. */
|
||||
writeVector(index: number, vector: number[]): void
|
||||
/** Write a flat batch starting at `startIndex`. Returns vectors written. */
|
||||
writeVectorsBatch(startIndex: number, vectorsFlat: number[]): number
|
||||
/** Read a single vector at `index`. Returns f64[]. */
|
||||
readVector(index: number): number[]
|
||||
/** Read multiple vectors by index. Returns flat f64[] of length n × dim. */
|
||||
readVectorsBatch(indices: number[]): number[]
|
||||
/** Prefetch via madvise(WILLNEED); no-op on unsupported platforms. */
|
||||
prefetch(indices: number[]): void
|
||||
/** Grow the vector file to a larger slot capacity. Re-maps the file. */
|
||||
resize(newCapacity: number): void
|
||||
/** msync — flush pending writes to disk. */
|
||||
flush(): void
|
||||
}
|
||||
|
||||
/**
|
||||
* Storage adapter factory — plugins register these to provide
|
||||
* new storage backends that users reference by name.
|
||||
|
|
|
|||
242
tests/unit/hnsw/mmap-vector-backend.test.ts
Normal file
242
tests/unit/hnsw/mmap-vector-backend.test.ts
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
/**
|
||||
* @module hnsw/mmapVectorBackend.test
|
||||
* @description Unit tests for the `MmapVectorBackend` bridge — the brainy-side
|
||||
* wrapper that translates UUID-keyed vector reads/writes into stable int slot
|
||||
* ops against an `vectorStore:mmap` provider.
|
||||
*
|
||||
* Mocks the provider so the tests run without cortex installed (cortex is a
|
||||
* downstream consumer of brainy, not a dev dep). The real integration with
|
||||
* cortex's `NativeMmapVectorStore` is exercised when cortex 2.4.0 picks up
|
||||
* this brainy release and re-runs its cross-language parity suite.
|
||||
*
|
||||
* Coverage:
|
||||
* 1. Open-then-write-then-read round-trips for a single vector.
|
||||
* 2. Batch reads return an array aligned with the input UUIDs, with `null`
|
||||
* entries for misses interleaved among hits — order preserved.
|
||||
* 3. Slot assignment is stable across multiple writes for the same UUID
|
||||
* (no re-slot, no overwrite of an adjacent slot).
|
||||
* 4. Writes beyond the initial capacity grow the file (doubling) without
|
||||
* losing the vectors already written.
|
||||
* 5. `readByUuid` returns `null` for both unknown UUIDs and UUIDs in the map
|
||||
* whose slot has not yet been written.
|
||||
* 6. Open is idempotent — opening an already-existing file reuses it.
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { mkdtemp, rm } from 'node:fs/promises'
|
||||
import { join } from 'node:path'
|
||||
import { MmapVectorBackend } from '../../../src/hnsw/mmapVectorBackend.js'
|
||||
import { EntityIdMapper } from '../../../src/utils/entityIdMapper.js'
|
||||
import type {
|
||||
VectorStoreMmapInstance,
|
||||
VectorStoreMmapProvider
|
||||
} from '../../../src/plugin.js'
|
||||
|
||||
/**
|
||||
* Minimal storage stub for the EntityIdMapper. The mapper only touches storage
|
||||
* in init/flush; for these tests the mapper starts empty and is never flushed.
|
||||
*/
|
||||
const stubStorage = {
|
||||
getMetadata: async () => undefined,
|
||||
saveMetadata: async () => {},
|
||||
getNouns: async () => ({ totalCount: 0, items: [] })
|
||||
} as any
|
||||
|
||||
/**
|
||||
* Pure in-memory mmap store. Mirrors cortex's NativeMmapVectorStore surface
|
||||
* just closely enough to exercise the backend's contract — no real mmap, no
|
||||
* file I/O, no f32 round-trip narrowing (the precision check is out of scope
|
||||
* here; that's covered by cortex's parity suite).
|
||||
*/
|
||||
class MockMmapStore implements VectorStoreMmapInstance {
|
||||
private readonly vectors: Array<number[] | undefined> = []
|
||||
private highestWritten = -1
|
||||
constructor(
|
||||
public readonly dim: number,
|
||||
private _capacity: number
|
||||
) {}
|
||||
get count(): number {
|
||||
return this.highestWritten + 1
|
||||
}
|
||||
get capacity(): number {
|
||||
return this._capacity
|
||||
}
|
||||
writeVector(index: number, vector: number[]): void {
|
||||
if (index >= this._capacity) {
|
||||
throw new Error(`Slot ${index} >= capacity ${this._capacity}`)
|
||||
}
|
||||
if (vector.length !== this.dim) {
|
||||
throw new Error(`Dim mismatch: expected ${this.dim}, got ${vector.length}`)
|
||||
}
|
||||
this.vectors[index] = [...vector]
|
||||
if (index > this.highestWritten) this.highestWritten = index
|
||||
}
|
||||
writeVectorsBatch(startIndex: number, vectorsFlat: number[]): number {
|
||||
if (vectorsFlat.length % this.dim !== 0) {
|
||||
throw new Error('vectorsFlat length not a multiple of dim')
|
||||
}
|
||||
const n = vectorsFlat.length / this.dim
|
||||
for (let i = 0; i < n; i++) {
|
||||
this.writeVector(startIndex + i, vectorsFlat.slice(i * this.dim, (i + 1) * this.dim))
|
||||
}
|
||||
return n
|
||||
}
|
||||
readVector(index: number): number[] {
|
||||
const v = this.vectors[index]
|
||||
if (!v) throw new Error(`Slot ${index} not written`)
|
||||
return [...v]
|
||||
}
|
||||
readVectorsBatch(indices: number[]): number[] {
|
||||
const flat: number[] = []
|
||||
for (const i of indices) {
|
||||
const v = this.vectors[i]
|
||||
if (!v) throw new Error(`Slot ${i} not written`)
|
||||
for (let k = 0; k < this.dim; k++) flat.push(v[k])
|
||||
}
|
||||
return flat
|
||||
}
|
||||
prefetch(_indices: number[]): void {
|
||||
/* no-op in mock */
|
||||
}
|
||||
resize(newCapacity: number): void {
|
||||
if (newCapacity < this._capacity) throw new Error('Cannot shrink')
|
||||
this._capacity = newCapacity
|
||||
}
|
||||
flush(): void {
|
||||
/* no-op in mock */
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Mock provider — keeps one MockMmapStore per path. open() throws if the path
|
||||
* doesn't exist yet (matches cortex semantics: open() is for existing files
|
||||
* only); create() throws if it does (cortex doesn't, but the brainy backend
|
||||
* does open-first-then-create, so the throw path is exercised).
|
||||
*/
|
||||
class MockMmapProvider implements VectorStoreMmapProvider {
|
||||
private files = new Map<string, MockMmapStore>()
|
||||
create(path: string, dim: number, capacity: number): VectorStoreMmapInstance {
|
||||
if (this.files.has(path)) throw new Error(`File exists at ${path}`)
|
||||
const store = new MockMmapStore(dim, capacity)
|
||||
this.files.set(path, store)
|
||||
return store
|
||||
}
|
||||
open(path: string): VectorStoreMmapInstance {
|
||||
const store = this.files.get(path)
|
||||
if (!store) throw new Error(`No file at ${path}`)
|
||||
return store
|
||||
}
|
||||
openReadOnly(path: string): VectorStoreMmapInstance {
|
||||
return this.open(path)
|
||||
}
|
||||
/** Test helper — peek at the underlying store. */
|
||||
_peek(path: string): MockMmapStore | undefined {
|
||||
return this.files.get(path)
|
||||
}
|
||||
}
|
||||
|
||||
describe('MmapVectorBackend (2.4.0 #2 — wraps vectorStore:mmap provider)', () => {
|
||||
let dir: string
|
||||
let path: string
|
||||
let idMapper: EntityIdMapper
|
||||
let provider: MockMmapProvider
|
||||
|
||||
beforeEach(async () => {
|
||||
dir = await mkdtemp(join(tmpdir(), 'brainy-mmap-vec-'))
|
||||
path = join(dir, 'vectors.bin')
|
||||
idMapper = new EntityIdMapper({ storage: stubStorage })
|
||||
await idMapper.init()
|
||||
provider = new MockMmapProvider()
|
||||
})
|
||||
|
||||
afterEach(async () => {
|
||||
await rm(dir, { recursive: true, force: true }).catch(() => {})
|
||||
})
|
||||
|
||||
it('open creates a new file when none exists, then round-trips a vector by UUID', async () => {
|
||||
const backend = await MmapVectorBackend.open(provider, path, 4, 16, idMapper)
|
||||
expect(backend.dim).toBe(4)
|
||||
|
||||
backend.writeByUuid('alpha', [1, 2, 3, 4])
|
||||
expect(backend.readByUuid('alpha')).toEqual([1, 2, 3, 4])
|
||||
})
|
||||
|
||||
it('reads return null for unknown UUIDs and for UUIDs in the map but not yet written', async () => {
|
||||
const backend = await MmapVectorBackend.open(provider, path, 2, 8, idMapper)
|
||||
|
||||
// Unknown — never seen by the mapper.
|
||||
expect(backend.readByUuid('ghost')).toBeNull()
|
||||
|
||||
// Known to the mapper but no slot ever written. We assign through the
|
||||
// mapper directly so the backend itself has not touched the slot.
|
||||
idMapper.getOrAssign('reserved')
|
||||
expect(backend.readByUuid('reserved')).toBeNull()
|
||||
})
|
||||
|
||||
it('batch read returns an array aligned to the input UUIDs (nulls preserved in place)', async () => {
|
||||
const backend = await MmapVectorBackend.open(provider, path, 3, 16, idMapper)
|
||||
backend.writeByUuid('a', [1, 1, 1])
|
||||
backend.writeByUuid('b', [2, 2, 2])
|
||||
backend.writeByUuid('c', [3, 3, 3])
|
||||
|
||||
// Interleave hits + a never-seen UUID + a duplicate hit. Order preserved.
|
||||
const result = backend.readBatchByUuid(['b', 'missing', 'a', 'c', 'missing-too', 'b'])
|
||||
expect(result).toEqual([
|
||||
[2, 2, 2],
|
||||
null,
|
||||
[1, 1, 1],
|
||||
[3, 3, 3],
|
||||
null,
|
||||
[2, 2, 2]
|
||||
])
|
||||
})
|
||||
|
||||
it('writes for the same UUID land in the same slot (stable id, no re-slotting)', async () => {
|
||||
const backend = await MmapVectorBackend.open(provider, path, 2, 8, idMapper)
|
||||
backend.writeByUuid('persistent', [1, 1])
|
||||
const slotAfterFirstWrite = idMapper.getInt('persistent')
|
||||
|
||||
backend.writeByUuid('persistent', [9, 9]) // overwrite the same slot
|
||||
expect(idMapper.getInt('persistent')).toBe(slotAfterFirstWrite)
|
||||
expect(backend.readByUuid('persistent')).toEqual([9, 9])
|
||||
|
||||
// Another UUID gets a different slot, and the first vector is undisturbed.
|
||||
backend.writeByUuid('other', [5, 5])
|
||||
expect(idMapper.getInt('other')).not.toBe(slotAfterFirstWrite)
|
||||
expect(backend.readByUuid('persistent')).toEqual([9, 9])
|
||||
expect(backend.readByUuid('other')).toEqual([5, 5])
|
||||
})
|
||||
|
||||
it('grows the file (doubling) when a write lands beyond capacity, without losing prior data', async () => {
|
||||
// Start at the smallest sane initial capacity (clamped to 16 by the backend).
|
||||
const backend = await MmapVectorBackend.open(provider, path, 2, 1, idMapper)
|
||||
const store = provider._peek(path)!
|
||||
expect(store.capacity).toBe(16) // backend floor
|
||||
|
||||
// Write 20 vectors so capacity must double at least once (16 → 32).
|
||||
const uuids: string[] = []
|
||||
for (let i = 0; i < 20; i++) {
|
||||
const uuid = `u-${i}`
|
||||
uuids.push(uuid)
|
||||
backend.writeByUuid(uuid, [i, i * 2])
|
||||
}
|
||||
expect(store.capacity).toBeGreaterThanOrEqual(32)
|
||||
|
||||
// Every vector survived the growth — no slot got overwritten or lost.
|
||||
for (let i = 0; i < uuids.length; i++) {
|
||||
expect(backend.readByUuid(uuids[i])).toEqual([i, i * 2])
|
||||
}
|
||||
})
|
||||
|
||||
it('opens an existing file idempotently when the second open hits the same path', async () => {
|
||||
// First open creates.
|
||||
const backend1 = await MmapVectorBackend.open(provider, path, 2, 8, idMapper)
|
||||
backend1.writeByUuid('persisted', [7, 7])
|
||||
|
||||
// Second open against the same path. The provider's create() throws on
|
||||
// collision; the backend's open-first behaviour means we reuse the file.
|
||||
const backend2 = await MmapVectorBackend.open(provider, path, 2, 8, idMapper)
|
||||
expect(backend2.readByUuid('persisted')).toEqual([7, 7])
|
||||
})
|
||||
})
|
||||
|
|
@ -22,7 +22,7 @@
|
|||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest'
|
||||
import { Brainy } from '../../src/brainy.js'
|
||||
import { Brainy } from '../../../src/brainy.js'
|
||||
|
||||
const DIM = 384
|
||||
const makeVec = (seed = 1) =>
|
||||
Loading…
Add table
Add a link
Reference in a new issue