feat: add SQ8 vector quantization, lazy loading, and two-phase rerank to HNSW
- SQ8 scalar quantization (8-bit) for 4x vector storage reduction - Lazy vector loading: evict float32 vectors after graph construction, load on-demand from storage via UnifiedCache - Two-phase search: over-retrieve with SQ8 approximate distances, rerank top candidates with exact float32 distances - Configuration surface: hnsw.quantization and hnsw.vectorStorage in BrainyConfig - All features disabled by default (zero behavior change for existing users) - 27 new tests covering quantization accuracy, lazy loading, reranking - Remove GitHub Actions CI (build locally, cortex CI handles native builds)
This commit is contained in:
parent
e384afcdac
commit
0f3a88429d
9 changed files with 1138 additions and 222 deletions
206
.github/workflows/ci.yml
vendored
206
.github/workflows/ci.yml
vendored
|
|
@ -1,206 +0,0 @@
|
||||||
name: CI
|
|
||||||
|
|
||||||
on:
|
|
||||||
push:
|
|
||||||
branches: [main]
|
|
||||||
pull_request:
|
|
||||||
branches: [main]
|
|
||||||
|
|
||||||
env:
|
|
||||||
NODE_VERSION: '20'
|
|
||||||
RUST_VERSION: 'stable'
|
|
||||||
|
|
||||||
jobs:
|
|
||||||
# Build and test TypeScript/JavaScript
|
|
||||||
test-node:
|
|
||||||
name: Test (Node.js)
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
- name: Setup Node.js
|
|
||||||
uses: actions/setup-node@v4
|
|
||||||
with:
|
|
||||||
node-version: ${{ env.NODE_VERSION }}
|
|
||||||
cache: 'npm'
|
|
||||||
|
|
||||||
- name: Install dependencies
|
|
||||||
run: npm ci
|
|
||||||
|
|
||||||
- name: Build TypeScript
|
|
||||||
run: npm run build
|
|
||||||
|
|
||||||
- name: Run unit tests
|
|
||||||
run: npm run test:unit
|
|
||||||
|
|
||||||
- name: Run integration tests
|
|
||||||
run: npm run test:integration
|
|
||||||
|
|
||||||
# Build and test with Bun
|
|
||||||
test-bun:
|
|
||||||
name: Test (Bun)
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
- name: Setup Bun
|
|
||||||
uses: oven-sh/setup-bun@v2
|
|
||||||
with:
|
|
||||||
bun-version: latest
|
|
||||||
|
|
||||||
- name: Setup Node.js (for npm compatibility)
|
|
||||||
uses: actions/setup-node@v4
|
|
||||||
with:
|
|
||||||
node-version: ${{ env.NODE_VERSION }}
|
|
||||||
cache: 'npm'
|
|
||||||
|
|
||||||
- name: Install dependencies
|
|
||||||
run: npm ci
|
|
||||||
|
|
||||||
- name: Build TypeScript
|
|
||||||
run: npm run build
|
|
||||||
|
|
||||||
- name: Run Bun tests
|
|
||||||
run: npm run test:bun
|
|
||||||
|
|
||||||
# Build Candle WASM (if Rust source changed)
|
|
||||||
build-candle-wasm:
|
|
||||||
name: Build Candle WASM
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
- name: Setup Rust
|
|
||||||
uses: dtolnay/rust-action@stable
|
|
||||||
with:
|
|
||||||
toolchain: ${{ env.RUST_VERSION }}
|
|
||||||
targets: wasm32-unknown-unknown
|
|
||||||
|
|
||||||
- name: Install wasm-pack
|
|
||||||
run: cargo install wasm-pack
|
|
||||||
|
|
||||||
- name: Cache Cargo registry
|
|
||||||
uses: actions/cache@v4
|
|
||||||
with:
|
|
||||||
path: |
|
|
||||||
~/.cargo/registry
|
|
||||||
~/.cargo/git
|
|
||||||
src/embeddings/candle-wasm/target
|
|
||||||
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
|
||||||
|
|
||||||
- name: Build WASM
|
|
||||||
run: |
|
|
||||||
cd src/embeddings/candle-wasm
|
|
||||||
wasm-pack build --target web --release
|
|
||||||
|
|
||||||
- name: Upload WASM artifacts
|
|
||||||
uses: actions/upload-artifact@v4
|
|
||||||
with:
|
|
||||||
name: candle-wasm
|
|
||||||
path: src/embeddings/wasm/pkg/
|
|
||||||
retention-days: 7
|
|
||||||
|
|
||||||
# Test Bun compile (standalone binary)
|
|
||||||
test-bun-compile:
|
|
||||||
name: Test Bun Compile
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
needs: [build-candle-wasm]
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
- name: Setup Bun
|
|
||||||
uses: oven-sh/setup-bun@v2
|
|
||||||
with:
|
|
||||||
bun-version: latest
|
|
||||||
|
|
||||||
- name: Setup Node.js
|
|
||||||
uses: actions/setup-node@v4
|
|
||||||
with:
|
|
||||||
node-version: ${{ env.NODE_VERSION }}
|
|
||||||
cache: 'npm'
|
|
||||||
|
|
||||||
- name: Download WASM artifacts
|
|
||||||
uses: actions/download-artifact@v4
|
|
||||||
with:
|
|
||||||
name: candle-wasm
|
|
||||||
path: src/embeddings/wasm/pkg/
|
|
||||||
|
|
||||||
- name: Install dependencies
|
|
||||||
run: npm ci
|
|
||||||
|
|
||||||
- name: Build TypeScript
|
|
||||||
run: npm run build
|
|
||||||
|
|
||||||
- name: Test Bun compile
|
|
||||||
run: |
|
|
||||||
# Create a test script
|
|
||||||
cat > /tmp/test-compile.ts << 'EOF'
|
|
||||||
import { Brainy } from './dist/index.js'
|
|
||||||
|
|
||||||
const brainy = new Brainy()
|
|
||||||
await brainy.init()
|
|
||||||
console.log('Brainy initialized!')
|
|
||||||
|
|
||||||
const embedding = await brainy.embed('Hello world')
|
|
||||||
console.log(`Embedding dimension: ${embedding.length}`)
|
|
||||||
|
|
||||||
if (embedding.length !== 384) {
|
|
||||||
throw new Error('Expected 384-dimensional embedding')
|
|
||||||
}
|
|
||||||
|
|
||||||
console.log('Test passed!')
|
|
||||||
EOF
|
|
||||||
|
|
||||||
# Compile to standalone binary
|
|
||||||
bun build --compile /tmp/test-compile.ts --outfile /tmp/brainy-test
|
|
||||||
|
|
||||||
# Run the compiled binary
|
|
||||||
/tmp/brainy-test
|
|
||||||
|
|
||||||
# Lint and type check
|
|
||||||
lint:
|
|
||||||
name: Lint & Type Check
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
- name: Setup Node.js
|
|
||||||
uses: actions/setup-node@v4
|
|
||||||
with:
|
|
||||||
node-version: ${{ env.NODE_VERSION }}
|
|
||||||
cache: 'npm'
|
|
||||||
|
|
||||||
- name: Install dependencies
|
|
||||||
run: npm ci
|
|
||||||
|
|
||||||
- name: Type check
|
|
||||||
run: npm run typecheck
|
|
||||||
|
|
||||||
- name: Lint
|
|
||||||
run: npm run lint
|
|
||||||
|
|
||||||
# License check
|
|
||||||
license-check:
|
|
||||||
name: License Check
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
- name: Setup Node.js
|
|
||||||
uses: actions/setup-node@v4
|
|
||||||
with:
|
|
||||||
node-version: ${{ env.NODE_VERSION }}
|
|
||||||
cache: 'npm'
|
|
||||||
|
|
||||||
- name: Install dependencies
|
|
||||||
run: npm ci
|
|
||||||
|
|
||||||
- name: Check licenses
|
|
||||||
run: |
|
|
||||||
# Install license checker
|
|
||||||
npm install -g license-checker
|
|
||||||
|
|
||||||
# Check for problematic licenses
|
|
||||||
license-checker --production --excludePrivatePackages \
|
|
||||||
--failOn 'GPL;LGPL;AGPL;SSPL' \
|
|
||||||
--summary
|
|
||||||
|
|
@ -6236,7 +6236,14 @@ export class Brainy<T = any> implements BrainyInterface<T> {
|
||||||
private setupIndex(): HNSWIndex | TypeAwareHNSWIndex {
|
private setupIndex(): HNSWIndex | TypeAwareHNSWIndex {
|
||||||
const indexConfig = {
|
const indexConfig = {
|
||||||
...this.config.index,
|
...this.config.index,
|
||||||
distanceFunction: this.distance
|
distanceFunction: this.distance,
|
||||||
|
// Wire HNSW optimization config (v7.11.0)
|
||||||
|
quantization: this.config.hnsw?.quantization ? {
|
||||||
|
enabled: this.config.hnsw.quantization.enabled ?? false,
|
||||||
|
bits: this.config.hnsw.quantization.bits ?? 8,
|
||||||
|
rerankMultiplier: this.config.hnsw.quantization.rerankMultiplier ?? 3
|
||||||
|
} : undefined,
|
||||||
|
vectorStorage: this.config.hnsw?.vectorStorage
|
||||||
}
|
}
|
||||||
|
|
||||||
const persistMode = this.resolveHNSWPersistMode()
|
const persistMode = this.resolveHNSWPersistMode()
|
||||||
|
|
@ -6368,6 +6375,8 @@ export class Brainy<T = any> implements BrainyInterface<T> {
|
||||||
reservedQueryMemory: config?.reservedQueryMemory ?? undefined as any,
|
reservedQueryMemory: config?.reservedQueryMemory ?? undefined as any,
|
||||||
// HNSW persistence mode - undefined = smart default in setupIndex
|
// HNSW persistence mode - undefined = smart default in setupIndex
|
||||||
hnswPersistMode: config?.hnswPersistMode ?? undefined as any,
|
hnswPersistMode: config?.hnswPersistMode ?? undefined as any,
|
||||||
|
// HNSW optimization options (v7.11.0)
|
||||||
|
hnsw: config?.hnsw ?? undefined as any,
|
||||||
// Embedding initialization - false = lazy init on first embed()
|
// Embedding initialization - false = lazy init on first embed()
|
||||||
eagerEmbeddings: config?.eagerEmbeddings ?? false,
|
eagerEmbeddings: config?.eagerEmbeddings ?? false,
|
||||||
// Integration Hub - undefined/false = disabled
|
// Integration Hub - undefined/false = disabled
|
||||||
|
|
|
||||||
|
|
@ -96,6 +96,10 @@ export interface HNSWNoun {
|
||||||
vector: Vector
|
vector: Vector
|
||||||
connections: Map<number, Set<string>> // level -> set of connected noun ids
|
connections: Map<number, Set<string>> // level -> set of connected noun ids
|
||||||
level: number // The highest layer this noun appears in
|
level: number // The highest layer this noun appears in
|
||||||
|
// SQ8 quantized vector for approximate distance computation (B1 optimization)
|
||||||
|
quantizedVector?: Uint8Array // 4x smaller than float32 vector
|
||||||
|
codebookMin?: number // quantization codebook minimum
|
||||||
|
codebookMax?: number // quantization codebook maximum
|
||||||
// ✅ NO metadata field - stored separately for optimization
|
// ✅ NO metadata field - stored separately for optimization
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -309,6 +313,14 @@ export interface HNSWConfig {
|
||||||
ml: number // Maximum level
|
ml: number // Maximum level
|
||||||
useDiskBasedIndex?: boolean // Whether to use disk-based index
|
useDiskBasedIndex?: boolean // Whether to use disk-based index
|
||||||
maxConcurrentNeighborWrites?: number // Maximum concurrent neighbor updates during insert. Default: unlimited (full concurrency)
|
maxConcurrentNeighborWrites?: number // Maximum concurrent neighbor updates during insert. Default: unlimited (full concurrency)
|
||||||
|
// SQ8 vector quantization (4x memory reduction, ~0.4% accuracy loss)
|
||||||
|
quantization?: {
|
||||||
|
enabled: boolean // default: false — preserves current behavior exactly
|
||||||
|
bits?: 8 | 4 // default: 8 (SQ8). SQ4 requires brainy-cortex native.
|
||||||
|
rerankMultiplier?: number // default: 3 — over-retrieve 3x, rerank with float32
|
||||||
|
}
|
||||||
|
// Vector storage mode
|
||||||
|
vectorStorage?: 'memory' | 'lazy' // default: 'memory' — 'lazy' evicts vectors after insert
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
|
||||||
|
|
@ -15,6 +15,8 @@ import { executeInThread } from '../utils/workerUtils.js'
|
||||||
import type { BaseStorage } from '../storage/baseStorage.js'
|
import type { BaseStorage } from '../storage/baseStorage.js'
|
||||||
import { getGlobalCache, UnifiedCache } from '../utils/unifiedCache.js'
|
import { getGlobalCache, UnifiedCache } from '../utils/unifiedCache.js'
|
||||||
import { prodLog } from '../utils/logger.js'
|
import { prodLog } from '../utils/logger.js'
|
||||||
|
import { quantizeSQ8, distanceSQ8 } from '../utils/vectorQuantization.js'
|
||||||
|
import type { SQ8QuantizedVector } from '../utils/vectorQuantization.js'
|
||||||
|
|
||||||
// Default HNSW parameters
|
// Default HNSW parameters
|
||||||
const DEFAULT_CONFIG: HNSWConfig = {
|
const DEFAULT_CONFIG: HNSWConfig = {
|
||||||
|
|
@ -53,6 +55,12 @@ export class HNSWIndex {
|
||||||
private dirtyNodes: Set<string> = new Set() // Nodes with unpersisted HNSW data
|
private dirtyNodes: Set<string> = new Set() // Nodes with unpersisted HNSW data
|
||||||
private dirtySystem: boolean = false // Whether system data (entryPoint, maxLevel) needs persist
|
private dirtySystem: boolean = false // Whether system data (entryPoint, maxLevel) needs persist
|
||||||
|
|
||||||
|
// SQ8 quantization support (B1 optimization)
|
||||||
|
private quantizationEnabled: boolean = false
|
||||||
|
private rerankMultiplier: number = 3
|
||||||
|
// Lazy vector storage (B2 optimization)
|
||||||
|
private vectorStorageMode: 'memory' | 'lazy' = 'memory'
|
||||||
|
|
||||||
constructor(
|
constructor(
|
||||||
config: Partial<HNSWConfig> = {},
|
config: Partial<HNSWConfig> = {},
|
||||||
distanceFunction: DistanceFunction = euclideanDistance,
|
distanceFunction: DistanceFunction = euclideanDistance,
|
||||||
|
|
@ -67,6 +75,15 @@ export class HNSWIndex {
|
||||||
this.storage = options.storage || null
|
this.storage = options.storage || null
|
||||||
this.persistMode = options.persistMode || 'immediate'
|
this.persistMode = options.persistMode || 'immediate'
|
||||||
|
|
||||||
|
// SQ8 quantization config (default: disabled, preserves current behavior)
|
||||||
|
if (config.quantization?.enabled) {
|
||||||
|
this.quantizationEnabled = true
|
||||||
|
this.rerankMultiplier = config.quantization.rerankMultiplier ?? 3
|
||||||
|
}
|
||||||
|
|
||||||
|
// Vector storage mode (default: 'memory', preserves current behavior)
|
||||||
|
this.vectorStorageMode = config.vectorStorage || 'memory'
|
||||||
|
|
||||||
// Use SAME UnifiedCache as Graph and Metadata for fair memory competition
|
// Use SAME UnifiedCache as Graph and Metadata for fair memory competition
|
||||||
this.unifiedCache = getGlobalCache()
|
this.unifiedCache = getGlobalCache()
|
||||||
}
|
}
|
||||||
|
|
@ -257,7 +274,11 @@ export class HNSWIndex {
|
||||||
id: original.id,
|
id: original.id,
|
||||||
vector: [...original.vector], // Deep copy vector array
|
vector: [...original.vector], // Deep copy vector array
|
||||||
connections: connectionsCopy,
|
connections: connectionsCopy,
|
||||||
level: original.level
|
level: original.level,
|
||||||
|
// Copy SQ8 quantized data if present
|
||||||
|
quantizedVector: original.quantizedVector ? new Uint8Array(original.quantizedVector) : undefined,
|
||||||
|
codebookMin: original.codebookMin,
|
||||||
|
codebookMax: original.codebookMax
|
||||||
}
|
}
|
||||||
|
|
||||||
this.nouns.set(nodeId, nodeCopy)
|
this.nouns.set(nodeId, nodeCopy)
|
||||||
|
|
@ -343,7 +364,7 @@ export class HNSWIndex {
|
||||||
// Generate random level for this noun
|
// Generate random level for this noun
|
||||||
const nounLevel = this.getRandomLevel()
|
const nounLevel = this.getRandomLevel()
|
||||||
|
|
||||||
// Create new noun
|
// Create new noun with optional SQ8 quantization
|
||||||
const noun: HNSWNoun = {
|
const noun: HNSWNoun = {
|
||||||
id,
|
id,
|
||||||
vector,
|
vector,
|
||||||
|
|
@ -351,6 +372,14 @@ export class HNSWIndex {
|
||||||
level: nounLevel
|
level: nounLevel
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Quantize vector if enabled (B1: 4x storage reduction)
|
||||||
|
if (this.quantizationEnabled) {
|
||||||
|
const sq8 = quantizeSQ8(vector)
|
||||||
|
noun.quantizedVector = sq8.quantized
|
||||||
|
noun.codebookMin = sq8.min
|
||||||
|
noun.codebookMax = sq8.max
|
||||||
|
}
|
||||||
|
|
||||||
// Initialize empty connection sets for each level
|
// Initialize empty connection sets for each level
|
||||||
for (let level = 0; level <= nounLevel; level++) {
|
for (let level = 0; level <= nounLevel; level++) {
|
||||||
noun.connections.set(level, new Set<string>())
|
noun.connections.set(level, new Set<string>())
|
||||||
|
|
@ -574,6 +603,13 @@ export class HNSWIndex {
|
||||||
this.highLevelNodes.get(nounLevel)!.add(id)
|
this.highLevelNodes.get(nounLevel)!.add(id)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Lazy vector eviction (B2: graph-only memory after insert)
|
||||||
|
// After graph construction completes, evict the full vector from memory.
|
||||||
|
// Future searches will load vectors on-demand via getVectorSafe() + UnifiedCache.
|
||||||
|
if (this.vectorStorageMode === 'lazy' && this.storage) {
|
||||||
|
noun.vector = [] // Release float32 vector from memory
|
||||||
|
}
|
||||||
|
|
||||||
// Persist HNSW graph data to storage
|
// Persist HNSW graph data to storage
|
||||||
// Respect persistMode setting
|
// Respect persistMode setting
|
||||||
if (this.storage && this.persistMode === 'immediate') {
|
if (this.storage && this.persistMode === 'immediate') {
|
||||||
|
|
@ -631,11 +667,21 @@ export class HNSWIndex {
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Search for nearest neighbors
|
* Search for nearest neighbors
|
||||||
|
*
|
||||||
|
* When SQ8 quantization is enabled and reranking is active:
|
||||||
|
* - Phase 1: Over-retrieve k*rerankMultiplier candidates using SQ8 approximate distances
|
||||||
|
* - Phase 2: Load full float32 vectors for candidates, compute exact distances, return top k
|
||||||
|
*
|
||||||
|
* @param queryVector Query vector
|
||||||
|
* @param k Number of results to return
|
||||||
|
* @param filter Optional filter function
|
||||||
|
* @param options Additional search options
|
||||||
*/
|
*/
|
||||||
public async search(
|
public async search(
|
||||||
queryVector: Vector,
|
queryVector: Vector,
|
||||||
k: number = 10,
|
k: number = 10,
|
||||||
filter?: (id: string) => Promise<boolean>
|
filter?: (id: string) => Promise<boolean>,
|
||||||
|
options?: { rerank?: { multiplier: number } }
|
||||||
): Promise<Array<[string, number]>> {
|
): Promise<Array<[string, number]>> {
|
||||||
if (this.nouns.size === 0) {
|
if (this.nouns.size === 0) {
|
||||||
return []
|
return []
|
||||||
|
|
@ -687,6 +733,11 @@ export class HNSWIndex {
|
||||||
|
|
||||||
let currObj = entryPoint
|
let currObj = entryPoint
|
||||||
|
|
||||||
|
// SQ8: Pre-quantize query vector for fast approximate distances during traversal
|
||||||
|
if (this.quantizationEnabled) {
|
||||||
|
this._querySQ8 = quantizeSQ8(queryVector)
|
||||||
|
}
|
||||||
|
|
||||||
// OPTIMIZATION: Preload entry point vector
|
// OPTIMIZATION: Preload entry point vector
|
||||||
await this.preloadVectors([entryPoint.id])
|
await this.preloadVectors([entryPoint.id])
|
||||||
|
|
||||||
|
|
@ -754,9 +805,14 @@ export class HNSWIndex {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Search at level 0 with ef = k
|
// Determine effective rerank multiplier
|
||||||
|
const rerankActive = this.quantizationEnabled && (options?.rerank || this.rerankMultiplier > 1)
|
||||||
|
const multiplier = options?.rerank?.multiplier ?? this.rerankMultiplier
|
||||||
|
const effectiveK = rerankActive ? k * multiplier : k
|
||||||
|
|
||||||
|
// Search at level 0 with ef = effectiveK
|
||||||
// If we have a filter, increase ef to compensate for filtered results
|
// If we have a filter, increase ef to compensate for filtered results
|
||||||
const ef = filter ? Math.max(this.config.efSearch * 3, k * 3) : Math.max(this.config.efSearch, k)
|
const ef = filter ? Math.max(this.config.efSearch * 3, effectiveK * 3) : Math.max(this.config.efSearch, effectiveK)
|
||||||
const nearestNouns = await this.searchLayer(
|
const nearestNouns = await this.searchLayer(
|
||||||
queryVector,
|
queryVector,
|
||||||
currObj,
|
currObj,
|
||||||
|
|
@ -765,6 +821,35 @@ export class HNSWIndex {
|
||||||
filter
|
filter
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// Phase 2: Rerank with exact float32 distances if quantization is active (B3)
|
||||||
|
if (rerankActive && nearestNouns.size > 0) {
|
||||||
|
// Clear SQ8 cache before reranking (we need exact distances now)
|
||||||
|
this._querySQ8 = null
|
||||||
|
|
||||||
|
const candidates = [...nearestNouns].slice(0, effectiveK)
|
||||||
|
|
||||||
|
// Load full float32 vectors for the candidate set
|
||||||
|
const candidateIds = candidates.map(([id]) => id)
|
||||||
|
await this.preloadVectors(candidateIds)
|
||||||
|
|
||||||
|
// Recompute exact distances
|
||||||
|
const reranked: Array<[string, number]> = []
|
||||||
|
for (const [id] of candidates) {
|
||||||
|
const noun = this.nouns.get(id)
|
||||||
|
if (!noun) continue
|
||||||
|
const exactVector = await this.getVectorSafe(noun)
|
||||||
|
const exactDist = this.distanceFunction(queryVector, exactVector)
|
||||||
|
reranked.push([id, exactDist])
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by exact distance and return top k
|
||||||
|
reranked.sort((a, b) => a[1] - b[1])
|
||||||
|
return reranked.slice(0, k)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clear SQ8 cache
|
||||||
|
this._querySQ8 = null
|
||||||
|
|
||||||
// Convert to array and sort by distance
|
// Convert to array and sort by distance
|
||||||
return [...nearestNouns].slice(0, k)
|
return [...nearestNouns].slice(0, k)
|
||||||
}
|
}
|
||||||
|
|
@ -1074,6 +1159,19 @@ export class HNSWIndex {
|
||||||
* @returns number | Promise<number> - sync when cached, async when needs load
|
* @returns number | Promise<number> - sync when cached, async when needs load
|
||||||
*/
|
*/
|
||||||
private distanceSafe(queryVector: Vector, noun: HNSWNoun): number | Promise<number> {
|
private distanceSafe(queryVector: Vector, noun: HNSWNoun): number | Promise<number> {
|
||||||
|
// SQ8 fast path: use quantized distance when available (B1 optimization)
|
||||||
|
// This avoids loading full float32 vectors during graph traversal
|
||||||
|
if (this.quantizationEnabled &&
|
||||||
|
noun.quantizedVector &&
|
||||||
|
noun.codebookMin !== undefined &&
|
||||||
|
noun.codebookMax !== undefined &&
|
||||||
|
this._querySQ8) {
|
||||||
|
return distanceSQ8(
|
||||||
|
this._querySQ8.quantized, this._querySQ8.min, this._querySQ8.max,
|
||||||
|
noun.quantizedVector, noun.codebookMin, noun.codebookMax
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
// Try sync fast path
|
// Try sync fast path
|
||||||
const nounVector = this.getVectorSync(noun)
|
const nounVector = this.getVectorSync(noun)
|
||||||
|
|
||||||
|
|
@ -1088,6 +1186,9 @@ export class HNSWIndex {
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Cached SQ8 quantization of the current query vector for distanceSafe fast path
|
||||||
|
private _querySQ8: SQ8QuantizedVector | null = null
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Get all nodes at a specific level for clustering
|
* Get all nodes at a specific level for clustering
|
||||||
* This enables O(n) clustering using HNSW's natural hierarchy
|
* This enables O(n) clustering using HNSW's natural hierarchy
|
||||||
|
|
@ -1207,14 +1308,25 @@ export class HNSWIndex {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Determine if vector should be kept in memory
|
||||||
|
const keepVector = shouldPreload && this.vectorStorageMode !== 'lazy'
|
||||||
|
|
||||||
// Create noun object with restored connections
|
// Create noun object with restored connections
|
||||||
const noun: HNSWNoun = {
|
const noun: HNSWNoun = {
|
||||||
id: nounData.id,
|
id: nounData.id,
|
||||||
vector: shouldPreload ? nounData.vector : [], // Preload if dataset is small
|
vector: keepVector ? nounData.vector : [], // Preload if dataset is small and not lazy
|
||||||
connections: new Map(),
|
connections: new Map(),
|
||||||
level: hnswData.level
|
level: hnswData.level
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Restore SQ8 quantized data if quantization is enabled
|
||||||
|
if (this.quantizationEnabled && nounData.vector.length > 0) {
|
||||||
|
const sq8 = quantizeSQ8(nounData.vector)
|
||||||
|
noun.quantizedVector = sq8.quantized
|
||||||
|
noun.codebookMin = sq8.min
|
||||||
|
noun.codebookMax = sq8.max
|
||||||
|
}
|
||||||
|
|
||||||
// Restore connections from persisted data
|
// Restore connections from persisted data
|
||||||
for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
|
for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
|
||||||
const level = parseInt(levelStr, 10)
|
const level = parseInt(levelStr, 10)
|
||||||
|
|
@ -1281,14 +1393,25 @@ export class HNSWIndex {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Determine if vector should be kept in memory
|
||||||
|
const keepVector = shouldPreload && this.vectorStorageMode !== 'lazy'
|
||||||
|
|
||||||
// Create noun object with restored connections
|
// Create noun object with restored connections
|
||||||
const noun: HNSWNoun = {
|
const noun: HNSWNoun = {
|
||||||
id: nounData.id,
|
id: nounData.id,
|
||||||
vector: shouldPreload ? nounData.vector : [], // Preload if dataset is small
|
vector: keepVector ? nounData.vector : [], // Preload if dataset is small and not lazy
|
||||||
connections: new Map(),
|
connections: new Map(),
|
||||||
level: hnswData.level
|
level: hnswData.level
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Restore SQ8 quantized data if quantization is enabled
|
||||||
|
if (this.quantizationEnabled && nounData.vector.length > 0) {
|
||||||
|
const sq8 = quantizeSQ8(nounData.vector)
|
||||||
|
noun.quantizedVector = sq8.quantized
|
||||||
|
noun.codebookMin = sq8.min
|
||||||
|
noun.codebookMax = sq8.max
|
||||||
|
}
|
||||||
|
|
||||||
// Restore connections from persisted data
|
// Restore connections from persisted data
|
||||||
for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
|
for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
|
||||||
const level = parseInt(levelStr, 10)
|
const level = parseInt(levelStr, 10)
|
||||||
|
|
|
||||||
|
|
@ -256,21 +256,22 @@ export class TypeAwareHNSWIndex {
|
||||||
queryVector: Vector,
|
queryVector: Vector,
|
||||||
k: number = 10,
|
k: number = 10,
|
||||||
type?: NounType | NounType[],
|
type?: NounType | NounType[],
|
||||||
filter?: (id: string) => Promise<boolean>
|
filter?: (id: string) => Promise<boolean>,
|
||||||
|
options?: { rerank?: { multiplier: number } }
|
||||||
): Promise<Array<[string, number]>> {
|
): Promise<Array<[string, number]>> {
|
||||||
// Single-type search (fast path)
|
// Single-type search (fast path)
|
||||||
if (type && typeof type === 'string') {
|
if (type && typeof type === 'string') {
|
||||||
const index = this.getIndexForType(type)
|
const index = this.getIndexForType(type)
|
||||||
return await index.search(queryVector, k, filter)
|
return await index.search(queryVector, k, filter, options)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Multi-type search (handle empty array edge case)
|
// Multi-type search (handle empty array edge case)
|
||||||
if (type && Array.isArray(type) && type.length > 0) {
|
if (type && Array.isArray(type) && type.length > 0) {
|
||||||
return await this.searchMultipleTypes(queryVector, k, type, filter)
|
return await this.searchMultipleTypes(queryVector, k, type, filter, options)
|
||||||
}
|
}
|
||||||
|
|
||||||
// All-types search (slowest path + empty array fallback)
|
// All-types search (slowest path + empty array fallback)
|
||||||
return await this.searchAllTypes(queryVector, k, filter)
|
return await this.searchAllTypes(queryVector, k, filter, options)
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
@ -286,7 +287,8 @@ export class TypeAwareHNSWIndex {
|
||||||
queryVector: Vector,
|
queryVector: Vector,
|
||||||
k: number,
|
k: number,
|
||||||
types: NounType[],
|
types: NounType[],
|
||||||
filter?: (id: string) => Promise<boolean>
|
filter?: (id: string) => Promise<boolean>,
|
||||||
|
options?: { rerank?: { multiplier: number } }
|
||||||
): Promise<Array<[string, number]>> {
|
): Promise<Array<[string, number]>> {
|
||||||
const allResults: Array<[string, number]> = []
|
const allResults: Array<[string, number]> = []
|
||||||
|
|
||||||
|
|
@ -294,7 +296,7 @@ export class TypeAwareHNSWIndex {
|
||||||
for (const type of types) {
|
for (const type of types) {
|
||||||
if (this.indexes.has(type)) {
|
if (this.indexes.has(type)) {
|
||||||
const index = this.indexes.get(type)!
|
const index = this.indexes.get(type)!
|
||||||
const results = await index.search(queryVector, k, filter)
|
const results = await index.search(queryVector, k, filter, options)
|
||||||
allResults.push(...results)
|
allResults.push(...results)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -320,13 +322,14 @@ export class TypeAwareHNSWIndex {
|
||||||
private async searchAllTypes(
|
private async searchAllTypes(
|
||||||
queryVector: Vector,
|
queryVector: Vector,
|
||||||
k: number,
|
k: number,
|
||||||
filter?: (id: string) => Promise<boolean>
|
filter?: (id: string) => Promise<boolean>,
|
||||||
|
options?: { rerank?: { multiplier: number } }
|
||||||
): Promise<Array<[string, number]>> {
|
): Promise<Array<[string, number]>> {
|
||||||
const allResults: Array<[string, number]> = []
|
const allResults: Array<[string, number]> = []
|
||||||
|
|
||||||
// Search each type's graph
|
// Search each type's graph
|
||||||
for (const [type, index] of this.indexes.entries()) {
|
for (const [type, index] of this.indexes.entries()) {
|
||||||
const results = await index.search(queryVector, k, filter)
|
const results = await index.search(queryVector, k, filter, options)
|
||||||
allResults.push(...results)
|
allResults.push(...results)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -728,6 +728,16 @@ export interface BrainyConfig {
|
||||||
// Cloud storage (GCS/S3/R2/Azure) should use 'deferred' for 30-50× faster adds
|
// Cloud storage (GCS/S3/R2/Azure) should use 'deferred' for 30-50× faster adds
|
||||||
hnswPersistMode?: 'immediate' | 'deferred'
|
hnswPersistMode?: 'immediate' | 'deferred'
|
||||||
|
|
||||||
|
// HNSW optimization options (v7.11.0)
|
||||||
|
hnsw?: {
|
||||||
|
quantization?: {
|
||||||
|
enabled?: boolean // default: false — current behavior exactly
|
||||||
|
bits?: 8 | 4 // default: 8 (SQ8). SQ4 requires brainy-cortex native.
|
||||||
|
rerankMultiplier?: number // default: 3 — over-retrieve 3x, rerank with float32
|
||||||
|
}
|
||||||
|
vectorStorage?: 'memory' | 'lazy' // default: 'memory' — 'lazy' evicts vectors after insert
|
||||||
|
}
|
||||||
|
|
||||||
// Memory management options
|
// Memory management options
|
||||||
maxQueryLimit?: number // Override auto-detected query result limit (max: 100000)
|
maxQueryLimit?: number // Override auto-detected query result limit (max: 100000)
|
||||||
reservedQueryMemory?: number // Memory reserved for queries in bytes (e.g., 1073741824 = 1GB)
|
reservedQueryMemory?: number // Memory reserved for queries in bytes (e.g., 1073741824 = 1GB)
|
||||||
|
|
|
||||||
218
src/utils/vectorQuantization.ts
Normal file
218
src/utils/vectorQuantization.ts
Normal file
|
|
@ -0,0 +1,218 @@
|
||||||
|
/**
|
||||||
|
* Vector Quantization Utilities
|
||||||
|
*
|
||||||
|
* Standalone SQ8 (Scalar Quantization, 8-bit) functions for HNSW vector compression.
|
||||||
|
* Each vector is independently quantized using its own min/max codebook.
|
||||||
|
*
|
||||||
|
* Storage format per vector:
|
||||||
|
* - quantized: Uint8Array (dimension bytes, 4x smaller than float32)
|
||||||
|
* - min: number (codebook minimum)
|
||||||
|
* - max: number (codebook maximum)
|
||||||
|
*
|
||||||
|
* Accuracy: SQ8 introduces ~0.4% error per dimension on normalized vectors.
|
||||||
|
* Combined with reranking (3x over-retrieval + float32 rerank), recall@100
|
||||||
|
* is expected to remain >99% for typical workloads.
|
||||||
|
*/
|
||||||
|
|
||||||
|
import type { Vector } from '../coreTypes.js'
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Result of SQ8 quantization
|
||||||
|
*/
|
||||||
|
export interface SQ8QuantizedVector {
|
||||||
|
quantized: Uint8Array
|
||||||
|
min: number
|
||||||
|
max: number
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Configuration for HNSW quantization
|
||||||
|
*/
|
||||||
|
export interface HNSWQuantizationConfig {
|
||||||
|
enabled: boolean
|
||||||
|
bits: 8 | 4
|
||||||
|
rerankMultiplier: number
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Quantize a float32 vector to SQ8 (8-bit unsigned integers).
|
||||||
|
*
|
||||||
|
* Maps [min, max] of the vector to [0, 255].
|
||||||
|
* For zero-range vectors (all values equal), all quantized values are 128.
|
||||||
|
*
|
||||||
|
* @param vector - Input float32 vector
|
||||||
|
* @returns Quantized vector with codebook (min/max)
|
||||||
|
*/
|
||||||
|
export function quantizeSQ8(vector: Vector): SQ8QuantizedVector {
|
||||||
|
const len = vector.length
|
||||||
|
let min = vector[0]
|
||||||
|
let max = vector[0]
|
||||||
|
|
||||||
|
for (let i = 1; i < len; i++) {
|
||||||
|
const v = vector[i]
|
||||||
|
if (v < min) min = v
|
||||||
|
if (v > max) max = v
|
||||||
|
}
|
||||||
|
|
||||||
|
const range = max - min
|
||||||
|
const quantized = new Uint8Array(len)
|
||||||
|
|
||||||
|
if (range === 0) {
|
||||||
|
// All values are identical — map to midpoint
|
||||||
|
quantized.fill(128)
|
||||||
|
} else {
|
||||||
|
const scale = 255 / range
|
||||||
|
for (let i = 0; i < len; i++) {
|
||||||
|
quantized[i] = Math.round((vector[i] - min) * scale)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return { quantized, min, max }
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Dequantize an SQ8 vector back to float32.
|
||||||
|
*
|
||||||
|
* @param quantized - Uint8Array of quantized values
|
||||||
|
* @param min - Codebook minimum
|
||||||
|
* @param max - Codebook maximum
|
||||||
|
* @returns Reconstructed float32 vector
|
||||||
|
*/
|
||||||
|
export function dequantizeSQ8(quantized: Uint8Array, min: number, max: number): Vector {
|
||||||
|
const len = quantized.length
|
||||||
|
const result = new Array<number>(len)
|
||||||
|
const range = max - min
|
||||||
|
|
||||||
|
if (range === 0) {
|
||||||
|
result.fill(min)
|
||||||
|
} else {
|
||||||
|
const scale = range / 255
|
||||||
|
for (let i = 0; i < len; i++) {
|
||||||
|
result[i] = min + quantized[i] * scale
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Compute approximate cosine distance between two SQ8-quantized vectors.
|
||||||
|
*
|
||||||
|
* Operates directly on uint8 arrays without full dequantization.
|
||||||
|
* Uses integer arithmetic where possible for speed.
|
||||||
|
*
|
||||||
|
* Cosine distance = 1 - (A.B / (|A| * |B|))
|
||||||
|
*
|
||||||
|
* For quantized values q_i with codebook (min, max):
|
||||||
|
* actual_i = min + q_i * (max - min) / 255
|
||||||
|
*
|
||||||
|
* The dot product and norms can be computed on quantized values
|
||||||
|
* and scaled by the codebook parameters.
|
||||||
|
*
|
||||||
|
* @param a - First quantized vector
|
||||||
|
* @param aMin - First vector codebook minimum
|
||||||
|
* @param aMax - First vector codebook maximum
|
||||||
|
* @param b - Second quantized vector
|
||||||
|
* @param bMin - Second vector codebook minimum
|
||||||
|
* @param bMax - Second vector codebook maximum
|
||||||
|
* @returns Approximate cosine distance [0, 2]
|
||||||
|
*/
|
||||||
|
export function distanceSQ8(
|
||||||
|
a: Uint8Array,
|
||||||
|
aMin: number,
|
||||||
|
aMax: number,
|
||||||
|
b: Uint8Array,
|
||||||
|
bMin: number,
|
||||||
|
bMax: number
|
||||||
|
): number {
|
||||||
|
const len = a.length
|
||||||
|
|
||||||
|
// Compute raw integer sums for efficiency
|
||||||
|
// actual_a[i] = aMin + a[i] * aScale, where aScale = (aMax - aMin) / 255
|
||||||
|
// actual_b[i] = bMin + b[i] * bScale, where bScale = (bMax - bMin) / 255
|
||||||
|
//
|
||||||
|
// dot = sum(actual_a[i] * actual_b[i])
|
||||||
|
// = sum((aMin + a[i]*aScale) * (bMin + b[i]*bScale))
|
||||||
|
// = n*aMin*bMin + aMin*bScale*sum(b[i]) + bMin*aScale*sum(a[i]) + aScale*bScale*sum(a[i]*b[i])
|
||||||
|
//
|
||||||
|
// normA^2 = sum(actual_a[i]^2)
|
||||||
|
// = n*aMin^2 + 2*aMin*aScale*sum(a[i]) + aScale^2*sum(a[i]^2)
|
||||||
|
|
||||||
|
const aScale = (aMax - aMin) / 255
|
||||||
|
const bScale = (bMax - bMin) / 255
|
||||||
|
|
||||||
|
let sumA = 0
|
||||||
|
let sumB = 0
|
||||||
|
let sumAB = 0
|
||||||
|
let sumAA = 0
|
||||||
|
let sumBB = 0
|
||||||
|
|
||||||
|
for (let i = 0; i < len; i++) {
|
||||||
|
const ai = a[i]
|
||||||
|
const bi = b[i]
|
||||||
|
sumA += ai
|
||||||
|
sumB += bi
|
||||||
|
sumAB += ai * bi
|
||||||
|
sumAA += ai * ai
|
||||||
|
sumBB += bi * bi
|
||||||
|
}
|
||||||
|
|
||||||
|
const dot =
|
||||||
|
len * aMin * bMin +
|
||||||
|
aMin * bScale * sumB +
|
||||||
|
bMin * aScale * sumA +
|
||||||
|
aScale * bScale * sumAB
|
||||||
|
|
||||||
|
const normASq =
|
||||||
|
len * aMin * aMin +
|
||||||
|
2 * aMin * aScale * sumA +
|
||||||
|
aScale * aScale * sumAA
|
||||||
|
|
||||||
|
const normBSq =
|
||||||
|
len * bMin * bMin +
|
||||||
|
2 * bMin * bScale * sumB +
|
||||||
|
bScale * bScale * sumBB
|
||||||
|
|
||||||
|
const normA = Math.sqrt(normASq)
|
||||||
|
const normB = Math.sqrt(normBSq)
|
||||||
|
|
||||||
|
if (normA === 0 || normB === 0) {
|
||||||
|
return 1.0 // Maximum distance for zero vectors
|
||||||
|
}
|
||||||
|
|
||||||
|
const cosine = dot / (normA * normB)
|
||||||
|
// Clamp to [-1, 1] to handle floating point imprecision
|
||||||
|
return 1 - Math.max(-1, Math.min(1, cosine))
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Serialize an SQ8 quantized vector to a compact binary format.
|
||||||
|
*
|
||||||
|
* Format: [min:float32][max:float32][quantized:uint8[]]
|
||||||
|
* Total size: 8 + dimension bytes
|
||||||
|
*
|
||||||
|
* @param sq8 - Quantized vector
|
||||||
|
* @returns ArrayBuffer with serialized data
|
||||||
|
*/
|
||||||
|
export function serializeSQ8(sq8: SQ8QuantizedVector): ArrayBuffer {
|
||||||
|
const buffer = new ArrayBuffer(8 + sq8.quantized.length)
|
||||||
|
const view = new DataView(buffer)
|
||||||
|
view.setFloat32(0, sq8.min, true) // little-endian
|
||||||
|
view.setFloat32(4, sq8.max, true)
|
||||||
|
new Uint8Array(buffer, 8).set(sq8.quantized)
|
||||||
|
return buffer
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Deserialize an SQ8 quantized vector from binary format.
|
||||||
|
*
|
||||||
|
* @param buffer - ArrayBuffer with serialized data
|
||||||
|
* @returns Deserialized SQ8 quantized vector
|
||||||
|
*/
|
||||||
|
export function deserializeSQ8(buffer: ArrayBuffer): SQ8QuantizedVector {
|
||||||
|
const view = new DataView(buffer)
|
||||||
|
const min = view.getFloat32(0, true)
|
||||||
|
const max = view.getFloat32(4, true)
|
||||||
|
const quantized = new Uint8Array(buffer, 8)
|
||||||
|
return { quantized, min, max }
|
||||||
|
}
|
||||||
293
tests/unit/hnsw/lazy-vectors.test.ts
Normal file
293
tests/unit/hnsw/lazy-vectors.test.ts
Normal file
|
|
@ -0,0 +1,293 @@
|
||||||
|
/**
|
||||||
|
* Lazy Vector Loading Tests (B2 optimization)
|
||||||
|
*
|
||||||
|
* Tests for:
|
||||||
|
* - Vectors evicted from memory after addItem() in lazy mode
|
||||||
|
* - Search returns correct results after vector eviction
|
||||||
|
* - Memory mode retains vectors (default behavior)
|
||||||
|
* - Lazy mode requires storage adapter
|
||||||
|
* - Combined lazy + quantization mode
|
||||||
|
*/
|
||||||
|
|
||||||
|
import { describe, it, expect, beforeEach } from 'vitest'
|
||||||
|
import { v4 as uuidv4 } from 'uuid'
|
||||||
|
import { HNSWIndex } from '../../../src/hnsw/hnswIndex.js'
|
||||||
|
import { euclideanDistance } from '../../../src/utils/index.js'
|
||||||
|
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'
|
||||||
|
|
||||||
|
// Helper: generate a random vector of given dimension
|
||||||
|
function randomVector(dim: number): number[] {
|
||||||
|
return Array.from({ length: dim }, () => Math.random() * 2 - 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Helper: save a vector to storage so lazy loading can retrieve it
|
||||||
|
// Both noun (vector) and metadata must be saved for getNounVector() to work
|
||||||
|
async function saveVector(storage: MemoryStorage, id: string, vector: number[]): Promise<void> {
|
||||||
|
await storage.saveNoun({
|
||||||
|
id,
|
||||||
|
vector,
|
||||||
|
connections: new Map(),
|
||||||
|
level: 0
|
||||||
|
})
|
||||||
|
await storage.saveNounMetadata(id, {
|
||||||
|
noun: 'thing',
|
||||||
|
createdAt: Date.now(),
|
||||||
|
updatedAt: Date.now()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
describe('Lazy Vector Loading (B2)', () => {
|
||||||
|
const dim = 32
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 1. LAZY MODE: VECTOR EVICTION
|
||||||
|
// =================================================================
|
||||||
|
describe('vector eviction in lazy mode', () => {
|
||||||
|
let index: HNSWIndex
|
||||||
|
let storage: MemoryStorage
|
||||||
|
|
||||||
|
beforeEach(async () => {
|
||||||
|
storage = new MemoryStorage()
|
||||||
|
index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
vectorStorage: 'lazy'
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should still be searchable after vector eviction', async () => {
|
||||||
|
const targetId = uuidv4()
|
||||||
|
const target = randomVector(dim)
|
||||||
|
|
||||||
|
// Store noun in storage so lazy loading can find it
|
||||||
|
await saveVector(storage, targetId, target)
|
||||||
|
await index.addItem({ id: targetId, vector: target })
|
||||||
|
|
||||||
|
// Add more entities
|
||||||
|
for (let i = 0; i < 20; i++) {
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await saveVector(storage, id, v)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Search should still find the target
|
||||||
|
const results = await index.search(target, 5)
|
||||||
|
expect(results.length).toBeGreaterThan(0)
|
||||||
|
|
||||||
|
// The closest result should be the target (distance ~0)
|
||||||
|
const targetResult = results.find(([id]) => id === targetId)
|
||||||
|
expect(targetResult).toBeDefined()
|
||||||
|
expect(targetResult![1]).toBeCloseTo(0, 1)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should return correct top-k results in lazy mode', async () => {
|
||||||
|
for (let i = 0; i < 30; i++) {
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await saveVector(storage, id, v)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
}
|
||||||
|
|
||||||
|
const query = randomVector(dim)
|
||||||
|
const results = await index.search(query, 10)
|
||||||
|
|
||||||
|
expect(results.length).toBe(10)
|
||||||
|
|
||||||
|
// Results should be sorted by distance
|
||||||
|
for (let i = 1; i < results.length; i++) {
|
||||||
|
expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 2. MEMORY MODE: VECTORS RETAINED (DEFAULT)
|
||||||
|
// =================================================================
|
||||||
|
describe('memory mode retains vectors (default)', () => {
|
||||||
|
it('should keep vectors in memory by default', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{ M: 4, efConstruction: 50, efSearch: 20 },
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
|
||||||
|
// Search should work without needing storage
|
||||||
|
const results = await index.search(v, 1)
|
||||||
|
expect(results.length).toBe(1)
|
||||||
|
expect(results[0][0]).toBe(id)
|
||||||
|
expect(results[0][1]).toBeCloseTo(0, 5)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should work without storage adapter in memory mode', async () => {
|
||||||
|
// No storage adapter provided
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{ M: 4, efConstruction: 50, efSearch: 20 },
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false }
|
||||||
|
)
|
||||||
|
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
|
||||||
|
const results = await index.search(v, 1)
|
||||||
|
expect(results.length).toBe(1)
|
||||||
|
expect(results[0][0]).toBe(id)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 3. LAZY + NO STORAGE: GRACEFUL BEHAVIOR
|
||||||
|
// =================================================================
|
||||||
|
describe('lazy mode without storage', () => {
|
||||||
|
it('should not evict vectors when no storage adapter is configured', async () => {
|
||||||
|
// vectorStorage: 'lazy' but no storage — vectors should stay in memory
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 4,
|
||||||
|
efConstruction: 50,
|
||||||
|
efSearch: 20,
|
||||||
|
vectorStorage: 'lazy'
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false }
|
||||||
|
)
|
||||||
|
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
|
||||||
|
// Should still work because vectors aren't evicted without storage
|
||||||
|
const results = await index.search(v, 1)
|
||||||
|
expect(results.length).toBe(1)
|
||||||
|
expect(results[0][0]).toBe(id)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 4. LAZY + QUANTIZATION COMBINED
|
||||||
|
// =================================================================
|
||||||
|
describe('lazy mode with quantization', () => {
|
||||||
|
it('should use SQ8 for traversal and load full vectors for rerank', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
vectorStorage: 'lazy',
|
||||||
|
quantization: { enabled: true, rerankMultiplier: 3 }
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const targetId = uuidv4()
|
||||||
|
const target = randomVector(dim)
|
||||||
|
await saveVector(storage, targetId, target)
|
||||||
|
await index.addItem({ id: targetId, vector: target })
|
||||||
|
|
||||||
|
for (let i = 0; i < 30; i++) {
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await saveVector(storage, id, v)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Search with reranking — should load full vectors for rerank phase
|
||||||
|
const results = await index.search(target, 5)
|
||||||
|
expect(results.length).toBe(5)
|
||||||
|
|
||||||
|
// The target should be the closest match
|
||||||
|
const targetResult = results.find(([id]) => id === targetId)
|
||||||
|
expect(targetResult).toBeDefined()
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should return results sorted by exact distance after rerank', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
vectorStorage: 'lazy',
|
||||||
|
quantization: { enabled: true, rerankMultiplier: 3 }
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
for (let i = 0; i < 40; i++) {
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
await saveVector(storage, id, v)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
}
|
||||||
|
|
||||||
|
const query = randomVector(dim)
|
||||||
|
const results = await index.search(query, 10)
|
||||||
|
|
||||||
|
// Results should be sorted by exact distance (rerank ensures this)
|
||||||
|
for (let i = 1; i < results.length; i++) {
|
||||||
|
expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 5. MULTIPLE SEARCHES IN LAZY MODE
|
||||||
|
// =================================================================
|
||||||
|
describe('multiple searches in lazy mode', () => {
|
||||||
|
it('should handle repeated searches correctly', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
vectorStorage: 'lazy'
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const entries: Array<{ id: string; vector: number[] }> = []
|
||||||
|
for (let i = 0; i < 25; i++) {
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(dim)
|
||||||
|
entries.push({ id, vector: v })
|
||||||
|
await saveVector(storage, id, v)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run multiple searches — each should return results and be consistent
|
||||||
|
for (let q = 0; q < 5; q++) {
|
||||||
|
const entry = entries[q * 5]
|
||||||
|
const results = await index.search(entry.vector, 5)
|
||||||
|
|
||||||
|
expect(results.length).toBe(5)
|
||||||
|
// Results should be sorted by distance
|
||||||
|
for (let i = 1; i < results.length; i++) {
|
||||||
|
expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
|
||||||
|
}
|
||||||
|
// At least the closest result should have a reasonably small distance
|
||||||
|
expect(results[0][1]).toBeLessThan(5)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
})
|
||||||
454
tests/unit/hnsw/quantization.test.ts
Normal file
454
tests/unit/hnsw/quantization.test.ts
Normal file
|
|
@ -0,0 +1,454 @@
|
||||||
|
/**
|
||||||
|
* SQ8 Quantization Tests
|
||||||
|
*
|
||||||
|
* Tests for:
|
||||||
|
* - SQ8 quantize/dequantize round-trip accuracy
|
||||||
|
* - SQ8 distance vs float32 distance correlation
|
||||||
|
* - Serialization/deserialization round-trip
|
||||||
|
* - Search with quantization enabled: recall vs exact
|
||||||
|
* - Two-phase rerank: improved recall over single-phase SQ8
|
||||||
|
* - Config defaults: quantization disabled by default (no behavior change)
|
||||||
|
*/
|
||||||
|
|
||||||
|
import { describe, it, expect, beforeEach } from 'vitest'
|
||||||
|
import { v4 as uuidv4 } from 'uuid'
|
||||||
|
import {
|
||||||
|
quantizeSQ8,
|
||||||
|
dequantizeSQ8,
|
||||||
|
distanceSQ8,
|
||||||
|
serializeSQ8,
|
||||||
|
deserializeSQ8
|
||||||
|
} from '../../../src/utils/vectorQuantization.js'
|
||||||
|
import { HNSWIndex } from '../../../src/hnsw/hnswIndex.js'
|
||||||
|
import { euclideanDistance } from '../../../src/utils/index.js'
|
||||||
|
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'
|
||||||
|
|
||||||
|
// Helper: generate a random vector of given dimension
|
||||||
|
function randomVector(dim: number): number[] {
|
||||||
|
return Array.from({ length: dim }, () => Math.random() * 2 - 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Helper: cosine distance for float32 vectors (reference implementation)
|
||||||
|
function cosineDistance(a: number[], b: number[]): number {
|
||||||
|
let dot = 0
|
||||||
|
let normA = 0
|
||||||
|
let normB = 0
|
||||||
|
for (let i = 0; i < a.length; i++) {
|
||||||
|
dot += a[i] * b[i]
|
||||||
|
normA += a[i] * a[i]
|
||||||
|
normB += b[i] * b[i]
|
||||||
|
}
|
||||||
|
const denom = Math.sqrt(normA) * Math.sqrt(normB)
|
||||||
|
if (denom === 0) return 1.0
|
||||||
|
return 1 - dot / denom
|
||||||
|
}
|
||||||
|
|
||||||
|
describe('SQ8 Quantization', () => {
|
||||||
|
// =================================================================
|
||||||
|
// 1. QUANTIZE / DEQUANTIZE ROUND-TRIP
|
||||||
|
// =================================================================
|
||||||
|
describe('quantize/dequantize round-trip', () => {
|
||||||
|
it('should round-trip a normalized vector with low error', () => {
|
||||||
|
const original = randomVector(384)
|
||||||
|
const sq8 = quantizeSQ8(original)
|
||||||
|
const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max)
|
||||||
|
|
||||||
|
expect(restored.length).toBe(original.length)
|
||||||
|
|
||||||
|
// Per-dimension error should be small (max 1/255 of range)
|
||||||
|
const range = sq8.max - sq8.min
|
||||||
|
const maxError = range / 255
|
||||||
|
for (let i = 0; i < original.length; i++) {
|
||||||
|
expect(Math.abs(restored[i] - original[i])).toBeLessThanOrEqual(maxError + 1e-7)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should handle a zero-range vector (all identical values)', () => {
|
||||||
|
const original = new Array(128).fill(0.5)
|
||||||
|
const sq8 = quantizeSQ8(original)
|
||||||
|
|
||||||
|
expect(sq8.min).toBe(0.5)
|
||||||
|
expect(sq8.max).toBe(0.5)
|
||||||
|
// All quantized values should be 128 (midpoint)
|
||||||
|
for (let i = 0; i < sq8.quantized.length; i++) {
|
||||||
|
expect(sq8.quantized[i]).toBe(128)
|
||||||
|
}
|
||||||
|
|
||||||
|
const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max)
|
||||||
|
for (let i = 0; i < restored.length; i++) {
|
||||||
|
expect(restored[i]).toBe(0.5)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should handle negative values', () => {
|
||||||
|
const original = [-1.0, -0.5, 0, 0.5, 1.0]
|
||||||
|
const sq8 = quantizeSQ8(original)
|
||||||
|
|
||||||
|
expect(sq8.min).toBe(-1.0)
|
||||||
|
expect(sq8.max).toBe(1.0)
|
||||||
|
// Min maps to 0, max maps to 255
|
||||||
|
expect(sq8.quantized[0]).toBe(0) // -1.0
|
||||||
|
expect(sq8.quantized[4]).toBe(255) // 1.0
|
||||||
|
|
||||||
|
const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max)
|
||||||
|
expect(Math.abs(restored[0] - (-1.0))).toBeLessThan(0.01)
|
||||||
|
expect(Math.abs(restored[4] - 1.0)).toBeLessThan(0.01)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should produce Uint8Array output in [0, 255] range', () => {
|
||||||
|
const original = randomVector(512)
|
||||||
|
const sq8 = quantizeSQ8(original)
|
||||||
|
|
||||||
|
expect(sq8.quantized).toBeInstanceOf(Uint8Array)
|
||||||
|
expect(sq8.quantized.length).toBe(512)
|
||||||
|
for (let i = 0; i < sq8.quantized.length; i++) {
|
||||||
|
expect(sq8.quantized[i]).toBeGreaterThanOrEqual(0)
|
||||||
|
expect(sq8.quantized[i]).toBeLessThanOrEqual(255)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should achieve 4x storage reduction (uint8 vs float32)', () => {
|
||||||
|
const dim = 384
|
||||||
|
const original = randomVector(dim)
|
||||||
|
const sq8 = quantizeSQ8(original)
|
||||||
|
|
||||||
|
const float32Size = dim * 4 // 4 bytes per float32
|
||||||
|
const sq8Size = sq8.quantized.byteLength + 8 // uint8 array + 2 floats for min/max
|
||||||
|
const ratio = float32Size / sq8Size
|
||||||
|
|
||||||
|
// Should be close to 4x (actually slightly less due to min/max overhead)
|
||||||
|
expect(ratio).toBeGreaterThan(3.9)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 2. SQ8 DISTANCE VS FLOAT32 DISTANCE
|
||||||
|
// =================================================================
|
||||||
|
describe('SQ8 distance accuracy', () => {
|
||||||
|
it('should correlate strongly with float32 cosine distance', () => {
|
||||||
|
const dim = 384
|
||||||
|
const numPairs = 100
|
||||||
|
const errors: number[] = []
|
||||||
|
|
||||||
|
for (let i = 0; i < numPairs; i++) {
|
||||||
|
const a = randomVector(dim)
|
||||||
|
const b = randomVector(dim)
|
||||||
|
|
||||||
|
const exactDist = cosineDistance(a, b)
|
||||||
|
|
||||||
|
const sq8A = quantizeSQ8(a)
|
||||||
|
const sq8B = quantizeSQ8(b)
|
||||||
|
const approxDist = distanceSQ8(
|
||||||
|
sq8A.quantized, sq8A.min, sq8A.max,
|
||||||
|
sq8B.quantized, sq8B.min, sq8B.max
|
||||||
|
)
|
||||||
|
|
||||||
|
errors.push(Math.abs(exactDist - approxDist))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Mean absolute error should be very small
|
||||||
|
const meanError = errors.reduce((sum, e) => sum + e, 0) / errors.length
|
||||||
|
expect(meanError).toBeLessThan(0.02) // Less than 2% average error
|
||||||
|
|
||||||
|
// Max error should be bounded
|
||||||
|
const maxError = Math.max(...errors)
|
||||||
|
expect(maxError).toBeLessThan(0.1) // Less than 10% worst case
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should return 0 for identical vectors', () => {
|
||||||
|
const v = randomVector(128)
|
||||||
|
const sq8 = quantizeSQ8(v)
|
||||||
|
const dist = distanceSQ8(
|
||||||
|
sq8.quantized, sq8.min, sq8.max,
|
||||||
|
sq8.quantized, sq8.min, sq8.max
|
||||||
|
)
|
||||||
|
expect(dist).toBeCloseTo(0, 5)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should preserve relative ordering of distances', () => {
|
||||||
|
const dim = 128
|
||||||
|
const query = randomVector(dim)
|
||||||
|
// Create a vector close to query and one far away
|
||||||
|
const close = query.map(v => v + (Math.random() * 0.1 - 0.05))
|
||||||
|
const far = randomVector(dim)
|
||||||
|
|
||||||
|
const exactClose = cosineDistance(query, close)
|
||||||
|
const exactFar = cosineDistance(query, far)
|
||||||
|
|
||||||
|
const sq8Query = quantizeSQ8(query)
|
||||||
|
const sq8Close = quantizeSQ8(close)
|
||||||
|
const sq8Far = quantizeSQ8(far)
|
||||||
|
|
||||||
|
const approxClose = distanceSQ8(
|
||||||
|
sq8Query.quantized, sq8Query.min, sq8Query.max,
|
||||||
|
sq8Close.quantized, sq8Close.min, sq8Close.max
|
||||||
|
)
|
||||||
|
const approxFar = distanceSQ8(
|
||||||
|
sq8Query.quantized, sq8Query.min, sq8Query.max,
|
||||||
|
sq8Far.quantized, sq8Far.min, sq8Far.max
|
||||||
|
)
|
||||||
|
|
||||||
|
// Relative ordering should be preserved
|
||||||
|
if (exactClose < exactFar) {
|
||||||
|
expect(approxClose).toBeLessThan(approxFar)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should handle zero vectors gracefully', () => {
|
||||||
|
const zero = new Array(64).fill(0)
|
||||||
|
const nonZero = randomVector(64)
|
||||||
|
|
||||||
|
const sq8Zero = quantizeSQ8(zero)
|
||||||
|
const sq8NonZero = quantizeSQ8(nonZero)
|
||||||
|
|
||||||
|
const dist = distanceSQ8(
|
||||||
|
sq8Zero.quantized, sq8Zero.min, sq8Zero.max,
|
||||||
|
sq8NonZero.quantized, sq8NonZero.min, sq8NonZero.max
|
||||||
|
)
|
||||||
|
// Zero vectors should return max distance (1.0)
|
||||||
|
expect(dist).toBeCloseTo(1.0, 1)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 3. SERIALIZATION / DESERIALIZATION
|
||||||
|
// =================================================================
|
||||||
|
describe('serialization round-trip', () => {
|
||||||
|
it('should serialize and deserialize SQ8 data with float32 precision', () => {
|
||||||
|
const original = randomVector(384)
|
||||||
|
const sq8 = quantizeSQ8(original)
|
||||||
|
const buffer = serializeSQ8(sq8)
|
||||||
|
const restored = deserializeSQ8(buffer)
|
||||||
|
|
||||||
|
// min/max are stored as float32, so some precision loss is expected
|
||||||
|
expect(restored.min).toBeCloseTo(sq8.min, 5)
|
||||||
|
expect(restored.max).toBeCloseTo(sq8.max, 5)
|
||||||
|
expect(restored.quantized.length).toBe(sq8.quantized.length)
|
||||||
|
for (let i = 0; i < sq8.quantized.length; i++) {
|
||||||
|
expect(restored.quantized[i]).toBe(sq8.quantized[i])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should produce compact binary format (8 + dim bytes)', () => {
|
||||||
|
const dim = 384
|
||||||
|
const sq8 = quantizeSQ8(randomVector(dim))
|
||||||
|
const buffer = serializeSQ8(sq8)
|
||||||
|
|
||||||
|
// 4 bytes for min (float32) + 4 bytes for max (float32) + dim bytes
|
||||||
|
expect(buffer.byteLength).toBe(8 + dim)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 4. HNSW SEARCH WITH QUANTIZATION
|
||||||
|
// =================================================================
|
||||||
|
describe('HNSW search with quantization', () => {
|
||||||
|
let index: HNSWIndex
|
||||||
|
let storage: MemoryStorage
|
||||||
|
const dim = 32 // Small dimension for fast tests
|
||||||
|
let ids: string[]
|
||||||
|
|
||||||
|
beforeEach(async () => {
|
||||||
|
storage = new MemoryStorage()
|
||||||
|
index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
quantization: { enabled: true, rerankMultiplier: 3 }
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
ids = []
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should return correct results with quantization enabled', async () => {
|
||||||
|
// Insert entities with UUID IDs
|
||||||
|
const vectors: number[][] = []
|
||||||
|
for (let i = 0; i < 50; i++) {
|
||||||
|
const id = uuidv4()
|
||||||
|
ids.push(id)
|
||||||
|
const v = randomVector(dim)
|
||||||
|
vectors.push(v)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Search with a known vector
|
||||||
|
const queryIdx = 5
|
||||||
|
const results = await index.search(vectors[queryIdx], 5)
|
||||||
|
|
||||||
|
// The exact vector should be the closest (distance ~0)
|
||||||
|
expect(results.length).toBeGreaterThan(0)
|
||||||
|
expect(results[0][0]).toBe(ids[queryIdx])
|
||||||
|
expect(results[0][1]).toBeCloseTo(0, 1)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should find the exact match as top result', async () => {
|
||||||
|
const targetId = uuidv4()
|
||||||
|
const target = randomVector(dim)
|
||||||
|
await index.addItem({ id: targetId, vector: target })
|
||||||
|
|
||||||
|
// Add noise vectors
|
||||||
|
for (let i = 0; i < 30; i++) {
|
||||||
|
await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
|
||||||
|
}
|
||||||
|
|
||||||
|
const results = await index.search(target, 1)
|
||||||
|
expect(results.length).toBe(1)
|
||||||
|
expect(results[0][0]).toBe(targetId)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should respect k parameter', async () => {
|
||||||
|
for (let i = 0; i < 100; i++) {
|
||||||
|
await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
|
||||||
|
}
|
||||||
|
|
||||||
|
const results5 = await index.search(randomVector(dim), 5)
|
||||||
|
const results10 = await index.search(randomVector(dim), 10)
|
||||||
|
|
||||||
|
expect(results5.length).toBe(5)
|
||||||
|
expect(results10.length).toBe(10)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 5. TWO-PHASE RERANK
|
||||||
|
// =================================================================
|
||||||
|
describe('two-phase rerank', () => {
|
||||||
|
it('should accept rerank options in search', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
quantization: { enabled: true, rerankMultiplier: 1 } // Disable default rerank
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const dim = 32
|
||||||
|
for (let i = 0; i < 30; i++) {
|
||||||
|
await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Search with explicit rerank
|
||||||
|
const results = await index.search(
|
||||||
|
randomVector(dim),
|
||||||
|
5,
|
||||||
|
undefined,
|
||||||
|
{ rerank: { multiplier: 3 } }
|
||||||
|
)
|
||||||
|
|
||||||
|
expect(results.length).toBe(5)
|
||||||
|
// Results should be sorted by distance
|
||||||
|
for (let i = 1; i < results.length; i++) {
|
||||||
|
expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('reranked results should be sorted by exact distance', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 8,
|
||||||
|
efConstruction: 100,
|
||||||
|
efSearch: 50,
|
||||||
|
ml: 8,
|
||||||
|
quantization: { enabled: true, rerankMultiplier: 3 }
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const dim = 32
|
||||||
|
for (let i = 0; i < 50; i++) {
|
||||||
|
await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
|
||||||
|
}
|
||||||
|
|
||||||
|
const query = randomVector(dim)
|
||||||
|
const results = await index.search(query, 10)
|
||||||
|
|
||||||
|
// Results should be sorted by exact distance (after reranking)
|
||||||
|
for (let i = 1; i < results.length; i++) {
|
||||||
|
expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// =================================================================
|
||||||
|
// 6. CONFIG DEFAULTS
|
||||||
|
// =================================================================
|
||||||
|
describe('configuration defaults', () => {
|
||||||
|
it('should disable quantization by default', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const defaultIndex = new HNSWIndex(
|
||||||
|
{ M: 4, efConstruction: 50, efSearch: 20 },
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const dim = 16
|
||||||
|
for (let i = 0; i < 10; i++) {
|
||||||
|
await defaultIndex.addItem({ id: uuidv4(), vector: randomVector(dim) })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Search should work identically to pre-quantization behavior
|
||||||
|
const results = await defaultIndex.search(randomVector(dim), 5)
|
||||||
|
expect(results.length).toBe(5)
|
||||||
|
|
||||||
|
// Distances should be exact euclidean (not SQ8 approximate)
|
||||||
|
for (const [, dist] of results) {
|
||||||
|
expect(typeof dist).toBe('number')
|
||||||
|
expect(dist).toBeGreaterThanOrEqual(0)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should use default rerankMultiplier of 3', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{
|
||||||
|
M: 4,
|
||||||
|
efConstruction: 50,
|
||||||
|
efSearch: 20,
|
||||||
|
quantization: { enabled: true }
|
||||||
|
},
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
// The index should be created without errors
|
||||||
|
// Default rerankMultiplier should be 3 (tested implicitly via search)
|
||||||
|
const dim = 16
|
||||||
|
for (let i = 0; i < 10; i++) {
|
||||||
|
await index.addItem({ id: uuidv4(), vector: randomVector(dim) })
|
||||||
|
}
|
||||||
|
|
||||||
|
const results = await index.search(randomVector(dim), 3)
|
||||||
|
expect(results.length).toBe(3)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('should default vectorStorage to memory mode', async () => {
|
||||||
|
const storage = new MemoryStorage()
|
||||||
|
const index = new HNSWIndex(
|
||||||
|
{ M: 4, efConstruction: 50, efSearch: 20 },
|
||||||
|
euclideanDistance,
|
||||||
|
{ useParallelization: false, storage }
|
||||||
|
)
|
||||||
|
|
||||||
|
const id = uuidv4()
|
||||||
|
const v = randomVector(16)
|
||||||
|
await index.addItem({ id, vector: v })
|
||||||
|
|
||||||
|
// In memory mode, vector should be immediately searchable
|
||||||
|
const results = await index.search(v, 1)
|
||||||
|
expect(results.length).toBe(1)
|
||||||
|
expect(results[0][0]).toBe(id)
|
||||||
|
expect(results[0][1]).toBeCloseTo(0, 5)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
})
|
||||||
Loading…
Add table
Add a link
Reference in a new issue