feat(8.0): zero-config finalize + cut JS quantization (config.vector = recall + persistMode)
This commit is contained in:
parent
f4c5d9749f
commit
f8e0079d3f
12 changed files with 348 additions and 1739 deletions
|
|
@ -6,7 +6,6 @@
|
|||
* - 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'
|
||||
|
|
@ -180,79 +179,7 @@ describe('Lazy Vector Loading (B2)', () => {
|
|||
})
|
||||
|
||||
// =================================================================
|
||||
// 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 JsHnswVectorIndex(
|
||||
{
|
||||
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 JsHnswVectorIndex(
|
||||
{
|
||||
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
|
||||
// 4. MULTIPLE SEARCHES IN LAZY MODE
|
||||
// =================================================================
|
||||
describe('multiple searches in lazy mode', () => {
|
||||
it('should handle repeated searches correctly', async () => {
|
||||
|
|
|
|||
|
|
@ -1,500 +0,0 @@
|
|||
/**
|
||||
* 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, afterEach } from 'vitest'
|
||||
import { v4 as uuidv4 } from 'uuid'
|
||||
import {
|
||||
quantizeSQ8,
|
||||
dequantizeSQ8,
|
||||
distanceSQ8,
|
||||
distanceSQ8Js,
|
||||
setSQ8DistanceImplementation,
|
||||
serializeSQ8,
|
||||
deserializeSQ8
|
||||
} from '../../../src/utils/vectorQuantization.js'
|
||||
import { JsHnswVectorIndex } 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)
|
||||
}
|
||||
|
||||
describe('SQ8 distance provider hook (setSQ8DistanceImplementation)', () => {
|
||||
const q1 = quantizeSQ8([0.1, 0.5, -0.3, 0.8])
|
||||
const q2 = quantizeSQ8([0.2, 0.4, -0.1, 0.7])
|
||||
|
||||
// Always restore the JS default so a swapped impl never leaks into other tests.
|
||||
afterEach(() => setSQ8DistanceImplementation(distanceSQ8Js))
|
||||
|
||||
it('dispatches to the JS implementation by default (identical result)', () => {
|
||||
const viaDispatcher = distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
|
||||
const viaJs = distanceSQ8Js(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
|
||||
expect(viaDispatcher).toBe(viaJs)
|
||||
expect(viaDispatcher).toBeGreaterThanOrEqual(0)
|
||||
expect(viaDispatcher).toBeLessThanOrEqual(2)
|
||||
})
|
||||
|
||||
it('routes through a swapped implementation with the cortex 6-arg signature', () => {
|
||||
const calls: Array<[Uint8Array, number, number, Uint8Array, number, number]> = []
|
||||
setSQ8DistanceImplementation((a, aMin, aMax, b, bMin, bMax) => {
|
||||
calls.push([a, aMin, aMax, b, bMin, bMax])
|
||||
return 0.4242
|
||||
})
|
||||
const result = distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
|
||||
expect(result).toBe(0.4242)
|
||||
expect(calls).toHaveLength(1)
|
||||
// Exactly the arguments cortex's native cosineDistanceSq8(a, aMin, aMax, b, bMin, bMax) expects.
|
||||
expect(calls[0][0]).toBe(q1.quantized)
|
||||
expect(calls[0][1]).toBe(q1.min)
|
||||
expect(calls[0][3]).toBe(q2.quantized)
|
||||
expect(calls[0][5]).toBe(q2.max)
|
||||
})
|
||||
|
||||
it('restores the JS default when set back', () => {
|
||||
setSQ8DistanceImplementation(() => -999)
|
||||
expect(distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)).toBe(-999)
|
||||
setSQ8DistanceImplementation(distanceSQ8Js)
|
||||
expect(distanceSQ8(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)).toBe(
|
||||
distanceSQ8Js(q1.quantized, q1.min, q1.max, q2.quantized, q2.min, q2.max)
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
// 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: JsHnswVectorIndex
|
||||
let storage: MemoryStorage
|
||||
const dim = 32 // Small dimension for fast tests
|
||||
let ids: string[]
|
||||
|
||||
beforeEach(async () => {
|
||||
storage = new MemoryStorage()
|
||||
index = new JsHnswVectorIndex(
|
||||
{
|
||||
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)
|
||||
// With quantization on a small graph (100 items), HNSW may occasionally
|
||||
// return slightly fewer than k due to approximation in distance calculations
|
||||
expect(results10.length).toBeGreaterThanOrEqual(8)
|
||||
expect(results10.length).toBeLessThanOrEqual(10)
|
||||
})
|
||||
})
|
||||
|
||||
// =================================================================
|
||||
// 5. TWO-PHASE RERANK
|
||||
// =================================================================
|
||||
describe('two-phase rerank', () => {
|
||||
it('should accept rerank options in search', async () => {
|
||||
const storage = new MemoryStorage()
|
||||
const index = new JsHnswVectorIndex(
|
||||
{
|
||||
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 JsHnswVectorIndex(
|
||||
{
|
||||
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 JsHnswVectorIndex(
|
||||
{ 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 JsHnswVectorIndex(
|
||||
{
|
||||
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 JsHnswVectorIndex(
|
||||
{ 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)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -117,10 +117,13 @@ describe('add() Performance Regression Tests', () => {
|
|||
expect(storageType).toBe('memory')
|
||||
})
|
||||
|
||||
it('should use immediate persistence for memory storage', async () => {
|
||||
// Memory storage should use immediate mode (already fast)
|
||||
it('should use deferred persistence for memory storage', async () => {
|
||||
// 8.0 adaptive default: memory storage uses 'deferred' — nothing survives
|
||||
// the process anyway, so per-add persistence writes are pure overhead.
|
||||
// (Filesystem storage gets 'immediate'; an explicit config.vector.persistMode
|
||||
// overrides either way.)
|
||||
const index = (brain as any).index
|
||||
expect(index.persistMode).toBe('immediate')
|
||||
expect(index.persistMode).toBe('deferred')
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -1,287 +0,0 @@
|
|||
/**
|
||||
* @module vectorQuantization-sq4.test
|
||||
* @description 2.5.0 #30 — SQ4 (4-bit scalar quantization).
|
||||
*
|
||||
* Pins down byte-for-byte correctness of the JS SQ4 reference implementation,
|
||||
* which must match cortex's Rust `quantize_sq4` / `dequantize_sq4` /
|
||||
* `cosine_distance_sq4` byte-for-byte (within float tolerance for the
|
||||
* distance — the integer accumulation is exact).
|
||||
*
|
||||
* Cross-language parity is the load-bearing invariant: brainy's HNSW SQ4
|
||||
* reranking swaps the JS distance for cortex's SIMD native via
|
||||
* `setSQ4DistanceImplementation`, and a recall drift between the two would
|
||||
* make search results differ depending on whether cortex is loaded.
|
||||
*
|
||||
* Coverage:
|
||||
* 1. Pack format — high nibble first, low nibble second; odd-dim trailing
|
||||
* pad nibble = 0. Total packed length = `ceil(dim / 2)`.
|
||||
* 2. Quantize → dequantize round-trip stays within the expected quantization
|
||||
* error envelope (each value is within `(max - min) / 15 / 2` of the
|
||||
* original — half a quantum step).
|
||||
* 3. Edge cases: empty vector throws; zero-range (all-equal) vectors map
|
||||
* every nibble to 8; even and odd dimensions both round-trip cleanly.
|
||||
* 4. Distance correctness: `distanceSQ4Js` on two quantized vectors agrees
|
||||
* with the true cosine distance on the dequantized values, within float
|
||||
* tolerance.
|
||||
* 5. Distance dispatch: `setSQ4DistanceImplementation` swaps the active fn;
|
||||
* `distanceSQ4` routes through the active implementation; restoring the
|
||||
* JS default reverts the swap.
|
||||
* 6. Serialization round-trip: `serializeSQ4` → `deserializeSQ4` preserves
|
||||
* all four fields (`quantized`, `min`, `max`, `dim`) byte-for-byte.
|
||||
*/
|
||||
|
||||
import { describe, it, expect, afterEach } from 'vitest'
|
||||
import {
|
||||
quantizeSQ4,
|
||||
dequantizeSQ4,
|
||||
distanceSQ4,
|
||||
distanceSQ4Js,
|
||||
setSQ4DistanceImplementation,
|
||||
serializeSQ4,
|
||||
deserializeSQ4
|
||||
} from '../../../src/utils/vectorQuantization.js'
|
||||
|
||||
/** True cosine distance on float vectors, for parity comparison against SQ4 distance. */
|
||||
function cosineDistanceFloat(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 - Math.max(-1, Math.min(1, dot / denom))
|
||||
}
|
||||
|
||||
/** Build a deterministic test vector with sine-modulated values across [-1, 1]. */
|
||||
function makeVec(dim: number, seed: number): number[] {
|
||||
const v = new Array<number>(dim)
|
||||
for (let i = 0; i < dim; i++) v[i] = Math.sin((i + seed) * 0.137)
|
||||
return v
|
||||
}
|
||||
|
||||
describe('SQ4 quantization (2.5.0 #30)', () => {
|
||||
// -----------------------------------------------------------------------
|
||||
// 1. Pack format — high nibble first, odd-dim trailing pad nibble = 0
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
describe('quantizeSQ4 — pack format', () => {
|
||||
it('packs two nibbles per byte, high nibble = vector[2i], low nibble = vector[2i+1]', () => {
|
||||
// Even dim. Choose a vector that lands on known nibbles:
|
||||
// min = 0, max = 15 → invRange = 1, so quantized[i] = round(vector[i]).
|
||||
const sq4 = quantizeSQ4([0, 15, 8, 1])
|
||||
expect(sq4.dim).toBe(4)
|
||||
expect(sq4.min).toBe(0)
|
||||
expect(sq4.max).toBe(15)
|
||||
expect(sq4.quantized.length).toBe(2)
|
||||
// byte 0 = (0 << 4) | 15 = 0x0f
|
||||
// byte 1 = (8 << 4) | 1 = 0x81
|
||||
expect(sq4.quantized[0]).toBe(0x0f)
|
||||
expect(sq4.quantized[1]).toBe(0x81)
|
||||
})
|
||||
|
||||
it('odd dim puts the final value in the high nibble and pads the low nibble with 0', () => {
|
||||
const sq4 = quantizeSQ4([0, 15, 7])
|
||||
expect(sq4.dim).toBe(3)
|
||||
expect(sq4.quantized.length).toBe(2)
|
||||
// byte 0 = (0 << 4) | 15 = 0x0f
|
||||
// byte 1 = (7 << 4) | 0 = 0x70 (low nibble is pad, 0)
|
||||
expect(sq4.quantized[0]).toBe(0x0f)
|
||||
expect(sq4.quantized[1]).toBe(0x70)
|
||||
})
|
||||
|
||||
it('packed length is ceil(dim / 2) for both even and odd dims', () => {
|
||||
expect(quantizeSQ4([0, 1]).quantized.length).toBe(1)
|
||||
expect(quantizeSQ4([0, 1, 2]).quantized.length).toBe(2)
|
||||
expect(quantizeSQ4([0, 1, 2, 3]).quantized.length).toBe(2)
|
||||
expect(quantizeSQ4([0, 1, 2, 3, 4]).quantized.length).toBe(3)
|
||||
expect(quantizeSQ4(makeVec(384, 1)).quantized.length).toBe(192)
|
||||
expect(quantizeSQ4(makeVec(385, 1)).quantized.length).toBe(193)
|
||||
})
|
||||
})
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 2. Quantize → dequantize round-trip is within the expected error envelope
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
describe('quantizeSQ4 / dequantizeSQ4 — round-trip error', () => {
|
||||
it('every reconstructed value is within half a quantum step of the original', () => {
|
||||
// The maximum reconstruction error for any value is half the quantum
|
||||
// step = (max - min) / 15 / 2. The factor 2 worst case happens at the
|
||||
// bin boundary; using 1.0 absolute tolerance plus the half-step bound
|
||||
// catches drift from the formula.
|
||||
const dims = [1, 2, 3, 4, 16, 17, 384, 385, 1000]
|
||||
for (const dim of dims) {
|
||||
const original = makeVec(dim, dim)
|
||||
const sq4 = quantizeSQ4(original)
|
||||
const restored = dequantizeSQ4(sq4.quantized, sq4.min, sq4.max, sq4.dim)
|
||||
expect(restored.length).toBe(dim)
|
||||
|
||||
const halfStep = (sq4.max - sq4.min) / 15 / 2
|
||||
for (let i = 0; i < dim; i++) {
|
||||
expect(Math.abs(restored[i] - original[i])).toBeLessThanOrEqual(halfStep + 1e-9)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
it('odd-dim vector loses NO data on round-trip — the pad nibble is discarded correctly', () => {
|
||||
const original = [-1.5, 0.0, 0.7, 1.3, 2.0]
|
||||
const sq4 = quantizeSQ4(original)
|
||||
const restored = dequantizeSQ4(sq4.quantized, sq4.min, sq4.max, sq4.dim)
|
||||
expect(restored.length).toBe(5)
|
||||
// The reconstruction has the same length and approximately the same
|
||||
// values; the trailing nibble pad does NOT appear in the output.
|
||||
})
|
||||
})
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 3. Edge cases
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
describe('edge cases', () => {
|
||||
it('zero-range vectors map every nibble to 8 (midpoint of [0, 15])', () => {
|
||||
const sq4 = quantizeSQ4([0.5, 0.5, 0.5, 0.5])
|
||||
expect(sq4.min).toBe(0.5)
|
||||
expect(sq4.max).toBe(0.5)
|
||||
expect(Array.from(sq4.quantized)).toEqual([0x88, 0x88])
|
||||
// Dequantize: range = 0, all values revert to min.
|
||||
const restored = dequantizeSQ4(sq4.quantized, sq4.min, sq4.max, sq4.dim)
|
||||
expect(restored).toEqual([0.5, 0.5, 0.5, 0.5])
|
||||
})
|
||||
|
||||
it('throws on empty input', () => {
|
||||
expect(() => quantizeSQ4([])).toThrow(/empty/i)
|
||||
})
|
||||
|
||||
it('clamps out-of-range integers to [0, 15] (defensive)', () => {
|
||||
// A single-value vector with max=min handled by zero-range branch;
|
||||
// here we use a wide range and confirm extremes land at the boundaries.
|
||||
const sq4 = quantizeSQ4([-100, 100, 0])
|
||||
// After mapping: min=-100, max=100, range=200. Quantized:
|
||||
// -100 → round(0 * 15) = 0
|
||||
// 100 → round(1 * 15) = 15
|
||||
// 0 → round(0.5 * 15) = 8
|
||||
// Packed: byte 0 = (0 << 4) | 15 = 0x0f, byte 1 = (8 << 4) | 0 = 0x80
|
||||
expect(Array.from(sq4.quantized)).toEqual([0x0f, 0x80])
|
||||
})
|
||||
})
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 4. Distance correctness — matches true cosine on dequantized values
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
describe('distanceSQ4Js — agreement with the dequantize-then-cosine baseline', () => {
|
||||
it('matches the true cosine distance on dequantized vectors within quantization error', () => {
|
||||
// For each pair, compute both the SQ4-direct distance and the
|
||||
// dequantize-then-cosine distance. They must agree within a tolerance
|
||||
// determined by quantization error — empirically a few percent on
|
||||
// 100-dim random vectors.
|
||||
const dim = 100
|
||||
for (let seed = 1; seed <= 10; seed++) {
|
||||
const a = makeVec(dim, seed)
|
||||
const b = makeVec(dim, seed + 50)
|
||||
const sa = quantizeSQ4(a)
|
||||
const sb = quantizeSQ4(b)
|
||||
|
||||
const sq4Dist = distanceSQ4Js(
|
||||
sa.quantized, sa.min, sa.max, sa.dim,
|
||||
sb.quantized, sb.min, sb.max, sb.dim
|
||||
)
|
||||
|
||||
const reA = dequantizeSQ4(sa.quantized, sa.min, sa.max, sa.dim)
|
||||
const reB = dequantizeSQ4(sb.quantized, sb.min, sb.max, sb.dim)
|
||||
const baselineDist = cosineDistanceFloat(reA, reB)
|
||||
|
||||
// The SQ4-direct path applies the EXACT same arithmetic as
|
||||
// cosine on dequantized values — the residual should be tiny float
|
||||
// accumulation error, never a percent-level drift.
|
||||
expect(Math.abs(sq4Dist - baselineDist)).toBeLessThan(1e-9)
|
||||
}
|
||||
})
|
||||
|
||||
it('throws on dimension mismatch', () => {
|
||||
const a = quantizeSQ4([0, 1, 2, 3])
|
||||
const b = quantizeSQ4([0, 1])
|
||||
expect(() => distanceSQ4Js(
|
||||
a.quantized, a.min, a.max, a.dim,
|
||||
b.quantized, b.min, b.max, b.dim
|
||||
)).toThrow(/dimension/i)
|
||||
})
|
||||
|
||||
it('returns 1.0 when either vector has zero norm (defensive)', () => {
|
||||
const z = quantizeSQ4([0, 0, 0, 0])
|
||||
const v = quantizeSQ4([1, 2, 3, 4])
|
||||
const d = distanceSQ4Js(
|
||||
z.quantized, z.min, z.max, z.dim,
|
||||
v.quantized, v.min, v.max, v.dim
|
||||
)
|
||||
// Zero-vector has range 0 → all-8 packed. SQ4 distance with the all-8
|
||||
// vector reconstructs to constant 0, so the formula yields cosine 0/0.
|
||||
// The reference impl returns 1.0 (maximum distance) for the degenerate case.
|
||||
expect(d).toBeGreaterThanOrEqual(0)
|
||||
expect(d).toBeLessThanOrEqual(2)
|
||||
})
|
||||
})
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 5. Distance dispatch — setSQ4DistanceImplementation hot-swap
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
describe('distance dispatch + setSQ4DistanceImplementation', () => {
|
||||
afterEach(() => {
|
||||
// Restore default after any swap.
|
||||
setSQ4DistanceImplementation(distanceSQ4Js)
|
||||
})
|
||||
|
||||
it('routes through the active implementation', () => {
|
||||
const a = quantizeSQ4(makeVec(64, 1))
|
||||
const b = quantizeSQ4(makeVec(64, 2))
|
||||
const baseline = distanceSQ4Js(a.quantized, a.min, a.max, a.dim, b.quantized, b.min, b.max, b.dim)
|
||||
const viaActive = distanceSQ4(a.quantized, a.min, a.max, a.dim, b.quantized, b.min, b.max, b.dim)
|
||||
expect(viaActive).toBe(baseline)
|
||||
})
|
||||
|
||||
it('setSQ4DistanceImplementation swaps the active fn; restoring reverts', () => {
|
||||
const sentinel = 0.4242
|
||||
setSQ4DistanceImplementation(() => sentinel)
|
||||
const a = quantizeSQ4([0, 1])
|
||||
const b = quantizeSQ4([1, 0])
|
||||
expect(distanceSQ4(a.quantized, a.min, a.max, a.dim, b.quantized, b.min, b.max, b.dim)).toBe(sentinel)
|
||||
|
||||
setSQ4DistanceImplementation(distanceSQ4Js)
|
||||
const restored = distanceSQ4(a.quantized, a.min, a.max, a.dim, b.quantized, b.min, b.max, b.dim)
|
||||
expect(restored).not.toBe(sentinel)
|
||||
})
|
||||
})
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 6. Serialization round-trip
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
describe('serialize / deserialize round-trip', () => {
|
||||
it('preserves quantized bytes, min, max, and dim byte-for-byte', () => {
|
||||
const original = quantizeSQ4(makeVec(384, 7))
|
||||
const buf = serializeSQ4(original)
|
||||
// 12-byte header (min:f32 + max:f32 + dim:u32 LE) + 192 packed bytes.
|
||||
expect(buf.byteLength).toBe(12 + 192)
|
||||
|
||||
const restored = deserializeSQ4(buf)
|
||||
expect(restored.dim).toBe(original.dim)
|
||||
// float32 round-trip on min/max may lose a few low bits (we stored as f32).
|
||||
expect(Math.abs(restored.min - original.min)).toBeLessThan(1e-5)
|
||||
expect(Math.abs(restored.max - original.max)).toBeLessThan(1e-5)
|
||||
expect(Array.from(restored.quantized)).toEqual(Array.from(original.quantized))
|
||||
})
|
||||
|
||||
it('works for odd dim (the trailing pad nibble round-trips correctly)', () => {
|
||||
const original = quantizeSQ4(makeVec(385, 11))
|
||||
const restored = deserializeSQ4(serializeSQ4(original))
|
||||
expect(restored.dim).toBe(385)
|
||||
expect(restored.quantized.length).toBe(193)
|
||||
expect(Array.from(restored.quantized)).toEqual(Array.from(original.quantized))
|
||||
})
|
||||
})
|
||||
})
|
||||
Loading…
Add table
Add a link
Reference in a new issue