David Snelling
f8c45f2d8d
Open source vector database with HNSW indexing, graph relationships, and metadata facets. Features CLI with professional augmentation registry integration for discovering extensions and capabilities.
471 lines
No EOL
17 KiB
JavaScript
471 lines
No EOL
17 KiB
JavaScript
/**
|
|
* Optimized HNSW (Hierarchical Navigable Small World) Index implementation
|
|
* Extends the base HNSW implementation with support for large datasets
|
|
* Uses product quantization for dimensionality reduction and disk-based storage when needed
|
|
*/
|
|
import { HNSWIndex } from './hnswIndex.js';
|
|
// Default configuration for the optimized HNSW index
|
|
const DEFAULT_OPTIMIZED_CONFIG = {
|
|
M: 16,
|
|
efConstruction: 200,
|
|
efSearch: 50,
|
|
ml: 16,
|
|
memoryThreshold: 1024 * 1024 * 1024, // 1GB default threshold
|
|
productQuantization: {
|
|
enabled: false,
|
|
numSubvectors: 16,
|
|
numCentroids: 256
|
|
},
|
|
useDiskBasedIndex: false
|
|
};
|
|
/**
|
|
* Product Quantization implementation
|
|
* Reduces vector dimensionality by splitting vectors into subvectors
|
|
* and quantizing each subvector to the nearest centroid
|
|
*/
|
|
class ProductQuantizer {
|
|
constructor(numSubvectors = 16, numCentroids = 256) {
|
|
this.centroids = [];
|
|
this.subvectorSize = 0;
|
|
this.initialized = false;
|
|
this.dimension = 0;
|
|
this.numSubvectors = numSubvectors;
|
|
this.numCentroids = numCentroids;
|
|
}
|
|
/**
|
|
* Initialize the product quantizer with training data
|
|
* @param vectors Training vectors to use for learning centroids
|
|
*/
|
|
train(vectors) {
|
|
if (vectors.length === 0) {
|
|
throw new Error('Cannot train product quantizer with empty vector set');
|
|
}
|
|
this.dimension = vectors[0].length;
|
|
this.subvectorSize = Math.ceil(this.dimension / this.numSubvectors);
|
|
// Initialize centroids for each subvector
|
|
for (let i = 0; i < this.numSubvectors; i++) {
|
|
// Extract subvectors from training data
|
|
const subvectors = vectors.map((vector) => {
|
|
const start = i * this.subvectorSize;
|
|
const end = Math.min(start + this.subvectorSize, this.dimension);
|
|
return vector.slice(start, end);
|
|
});
|
|
// Initialize centroids for this subvector using k-means++
|
|
this.centroids[i] = this.kMeansPlusPlus(subvectors, this.numCentroids);
|
|
}
|
|
this.initialized = true;
|
|
}
|
|
/**
|
|
* Quantize a vector using product quantization
|
|
* @param vector Vector to quantize
|
|
* @returns Array of centroid indices, one for each subvector
|
|
*/
|
|
quantize(vector) {
|
|
if (!this.initialized) {
|
|
throw new Error('Product quantizer not initialized. Call train() first.');
|
|
}
|
|
if (vector.length !== this.dimension) {
|
|
throw new Error(`Vector dimension mismatch: expected ${this.dimension}, got ${vector.length}`);
|
|
}
|
|
const codes = [];
|
|
// Quantize each subvector
|
|
for (let i = 0; i < this.numSubvectors; i++) {
|
|
const start = i * this.subvectorSize;
|
|
const end = Math.min(start + this.subvectorSize, this.dimension);
|
|
const subvector = vector.slice(start, end);
|
|
// Find nearest centroid
|
|
let minDist = Number.MAX_VALUE;
|
|
let nearestCentroidIndex = 0;
|
|
for (let j = 0; j < this.centroids[i].length; j++) {
|
|
const centroid = this.centroids[i][j];
|
|
const dist = this.euclideanDistanceSquared(subvector, centroid);
|
|
if (dist < minDist) {
|
|
minDist = dist;
|
|
nearestCentroidIndex = j;
|
|
}
|
|
}
|
|
codes.push(nearestCentroidIndex);
|
|
}
|
|
return codes;
|
|
}
|
|
/**
|
|
* Reconstruct a vector from its quantized representation
|
|
* @param codes Array of centroid indices
|
|
* @returns Reconstructed vector
|
|
*/
|
|
reconstruct(codes) {
|
|
if (!this.initialized) {
|
|
throw new Error('Product quantizer not initialized. Call train() first.');
|
|
}
|
|
if (codes.length !== this.numSubvectors) {
|
|
throw new Error(`Code length mismatch: expected ${this.numSubvectors}, got ${codes.length}`);
|
|
}
|
|
const reconstructed = [];
|
|
// Reconstruct each subvector
|
|
for (let i = 0; i < this.numSubvectors; i++) {
|
|
const centroidIndex = codes[i];
|
|
const centroid = this.centroids[i][centroidIndex];
|
|
// Add centroid components to reconstructed vector
|
|
for (const component of centroid) {
|
|
reconstructed.push(component);
|
|
}
|
|
}
|
|
// Trim to original dimension if needed
|
|
return reconstructed.slice(0, this.dimension);
|
|
}
|
|
/**
|
|
* Compute squared Euclidean distance between two vectors
|
|
* @param a First vector
|
|
* @param b Second vector
|
|
* @returns Squared Euclidean distance
|
|
*/
|
|
euclideanDistanceSquared(a, b) {
|
|
let sum = 0;
|
|
const length = Math.min(a.length, b.length);
|
|
for (let i = 0; i < length; i++) {
|
|
const diff = a[i] - b[i];
|
|
sum += diff * diff;
|
|
}
|
|
return sum;
|
|
}
|
|
/**
|
|
* Implement k-means++ algorithm to initialize centroids
|
|
* @param vectors Vectors to cluster
|
|
* @param k Number of clusters
|
|
* @returns Array of centroids
|
|
*/
|
|
kMeansPlusPlus(vectors, k) {
|
|
if (vectors.length < k) {
|
|
// If we have fewer vectors than centroids, use the vectors as centroids
|
|
return [...vectors];
|
|
}
|
|
const centroids = [];
|
|
// Choose first centroid randomly
|
|
const firstIndex = Math.floor(Math.random() * vectors.length);
|
|
centroids.push([...vectors[firstIndex]]);
|
|
// Choose remaining centroids
|
|
for (let i = 1; i < k; i++) {
|
|
// Compute distances to nearest centroid for each vector
|
|
const distances = vectors.map((vector) => {
|
|
let minDist = Number.MAX_VALUE;
|
|
for (const centroid of centroids) {
|
|
const dist = this.euclideanDistanceSquared(vector, centroid);
|
|
minDist = Math.min(minDist, dist);
|
|
}
|
|
return minDist;
|
|
});
|
|
// Compute sum of distances
|
|
const distSum = distances.reduce((sum, dist) => sum + dist, 0);
|
|
// Choose next centroid with probability proportional to distance
|
|
let r = Math.random() * distSum;
|
|
let nextIndex = 0;
|
|
for (let j = 0; j < distances.length; j++) {
|
|
r -= distances[j];
|
|
if (r <= 0) {
|
|
nextIndex = j;
|
|
break;
|
|
}
|
|
}
|
|
centroids.push([...vectors[nextIndex]]);
|
|
}
|
|
return centroids;
|
|
}
|
|
/**
|
|
* Get the centroids for each subvector
|
|
* @returns Array of centroid arrays
|
|
*/
|
|
getCentroids() {
|
|
return this.centroids;
|
|
}
|
|
/**
|
|
* Set the centroids for each subvector
|
|
* @param centroids Array of centroid arrays
|
|
*/
|
|
setCentroids(centroids) {
|
|
this.centroids = centroids;
|
|
this.numSubvectors = centroids.length;
|
|
this.numCentroids = centroids[0].length;
|
|
this.initialized = true;
|
|
}
|
|
/**
|
|
* Get the dimension of the vectors
|
|
* @returns Dimension
|
|
*/
|
|
getDimension() {
|
|
return this.dimension;
|
|
}
|
|
/**
|
|
* Set the dimension of the vectors
|
|
* @param dimension Dimension
|
|
*/
|
|
setDimension(dimension) {
|
|
this.dimension = dimension;
|
|
this.subvectorSize = Math.ceil(dimension / this.numSubvectors);
|
|
}
|
|
}
|
|
/**
|
|
* Optimized HNSW Index implementation
|
|
* Extends the base HNSW implementation with support for large datasets
|
|
* Uses product quantization for dimensionality reduction and disk-based storage when needed
|
|
*/
|
|
export class HNSWIndexOptimized extends HNSWIndex {
|
|
constructor(config = {}, distanceFunction, storage = null) {
|
|
// Initialize base HNSW index with standard config
|
|
super(config, distanceFunction);
|
|
this.productQuantizer = null;
|
|
this.storage = null;
|
|
this.useDiskBasedIndex = false;
|
|
this.useProductQuantization = false;
|
|
this.quantizedVectors = new Map();
|
|
this.memoryUsage = 0;
|
|
this.vectorCount = 0;
|
|
// Thread safety for memory usage tracking
|
|
this.memoryUpdateLock = Promise.resolve();
|
|
// Set optimized config
|
|
this.optimizedConfig = { ...DEFAULT_OPTIMIZED_CONFIG, ...config };
|
|
// Set storage adapter
|
|
this.storage = storage;
|
|
// Initialize product quantizer if enabled
|
|
if (this.optimizedConfig.productQuantization?.enabled) {
|
|
this.useProductQuantization = true;
|
|
this.productQuantizer = new ProductQuantizer(this.optimizedConfig.productQuantization.numSubvectors, this.optimizedConfig.productQuantization.numCentroids);
|
|
}
|
|
// Set disk-based index flag
|
|
this.useDiskBasedIndex = this.optimizedConfig.useDiskBasedIndex || false;
|
|
}
|
|
/**
|
|
* Thread-safe method to update memory usage
|
|
* @param memoryDelta Change in memory usage (can be negative)
|
|
* @param vectorCountDelta Change in vector count (can be negative)
|
|
*/
|
|
async updateMemoryUsage(memoryDelta, vectorCountDelta) {
|
|
this.memoryUpdateLock = this.memoryUpdateLock.then(async () => {
|
|
this.memoryUsage = Math.max(0, this.memoryUsage + memoryDelta);
|
|
this.vectorCount = Math.max(0, this.vectorCount + vectorCountDelta);
|
|
});
|
|
await this.memoryUpdateLock;
|
|
}
|
|
/**
|
|
* Thread-safe method to get current memory usage
|
|
* @returns Current memory usage and vector count
|
|
*/
|
|
async getMemoryUsageAsync() {
|
|
await this.memoryUpdateLock;
|
|
return {
|
|
memoryUsage: this.memoryUsage,
|
|
vectorCount: this.vectorCount
|
|
};
|
|
}
|
|
/**
|
|
* Add a vector to the index
|
|
* Uses product quantization if enabled and memory threshold is exceeded
|
|
*/
|
|
async addItem(item) {
|
|
// Check if item is defined
|
|
if (!item) {
|
|
throw new Error('Item is undefined or null');
|
|
}
|
|
const { id, vector } = item;
|
|
// Check if vector is defined
|
|
if (!vector) {
|
|
throw new Error('Vector is undefined or null');
|
|
}
|
|
// Estimate memory usage for this vector
|
|
const vectorMemory = vector.length * 8; // 8 bytes per number (Float64)
|
|
const connectionsMemory = this.optimizedConfig.M * this.optimizedConfig.ml * 16; // Estimate for connections
|
|
const totalMemory = vectorMemory + connectionsMemory;
|
|
// Update memory usage estimate (thread-safe)
|
|
await this.updateMemoryUsage(totalMemory, 1);
|
|
// Check if we should switch to product quantization
|
|
const currentMemoryUsage = await this.getMemoryUsageAsync();
|
|
if (this.useProductQuantization &&
|
|
currentMemoryUsage.memoryUsage > this.optimizedConfig.memoryThreshold &&
|
|
this.productQuantizer &&
|
|
!this.productQuantizer.getDimension()) {
|
|
// Initialize product quantizer with existing vectors
|
|
this.initializeProductQuantizer();
|
|
}
|
|
// If product quantization is active, quantize the vector
|
|
if (this.useProductQuantization &&
|
|
this.productQuantizer &&
|
|
this.productQuantizer.getDimension() > 0) {
|
|
// Quantize the vector
|
|
const codes = this.productQuantizer.quantize(vector);
|
|
// Store the quantized vector
|
|
this.quantizedVectors.set(id, codes);
|
|
// Reconstruct the vector for indexing
|
|
const reconstructedVector = this.productQuantizer.reconstruct(codes);
|
|
// Add the reconstructed vector to the index
|
|
return await super.addItem({ id, vector: reconstructedVector });
|
|
}
|
|
// If disk-based index is active and storage is available, store the vector
|
|
if (this.useDiskBasedIndex && this.storage) {
|
|
// Create a noun object
|
|
const noun = {
|
|
id,
|
|
vector,
|
|
connections: new Map(),
|
|
level: 0
|
|
};
|
|
// Store the noun
|
|
this.storage.saveNoun(noun).catch((error) => {
|
|
console.error(`Failed to save noun ${id} to storage:`, error);
|
|
});
|
|
}
|
|
// Add the vector to the in-memory index
|
|
return await super.addItem(item);
|
|
}
|
|
/**
|
|
* Search for nearest neighbors
|
|
* Uses product quantization if enabled
|
|
*/
|
|
async search(queryVector, k = 10) {
|
|
// Check if query vector is defined
|
|
if (!queryVector) {
|
|
throw new Error('Query vector is undefined or null');
|
|
}
|
|
// If product quantization is active, quantize the query vector
|
|
if (this.useProductQuantization &&
|
|
this.productQuantizer &&
|
|
this.productQuantizer.getDimension() > 0) {
|
|
// Quantize the query vector
|
|
const codes = this.productQuantizer.quantize(queryVector);
|
|
// Reconstruct the query vector
|
|
const reconstructedVector = this.productQuantizer.reconstruct(codes);
|
|
// Search with the reconstructed vector
|
|
return await super.search(reconstructedVector, k);
|
|
}
|
|
// Otherwise, use the standard search
|
|
return await super.search(queryVector, k);
|
|
}
|
|
/**
|
|
* Remove an item from the index
|
|
*/
|
|
removeItem(id) {
|
|
// If product quantization is active, remove the quantized vector
|
|
if (this.useProductQuantization) {
|
|
this.quantizedVectors.delete(id);
|
|
}
|
|
// If disk-based index is active and storage is available, remove the vector from storage
|
|
if (this.useDiskBasedIndex && this.storage) {
|
|
this.storage.deleteNoun(id).catch((error) => {
|
|
console.error(`Failed to delete noun ${id} from storage:`, error);
|
|
});
|
|
}
|
|
// Update memory usage estimate (async operation, but don't block removal)
|
|
this.getMemoryUsageAsync().then((currentMemoryUsage) => {
|
|
if (currentMemoryUsage.vectorCount > 0) {
|
|
const memoryPerVector = currentMemoryUsage.memoryUsage / currentMemoryUsage.vectorCount;
|
|
this.updateMemoryUsage(-memoryPerVector, -1);
|
|
}
|
|
}).catch((error) => {
|
|
console.error('Failed to update memory usage after removal:', error);
|
|
});
|
|
// Remove the item from the in-memory index
|
|
return super.removeItem(id);
|
|
}
|
|
/**
|
|
* Clear the index
|
|
*/
|
|
async clear() {
|
|
// Clear product quantization data
|
|
if (this.useProductQuantization) {
|
|
this.quantizedVectors.clear();
|
|
this.productQuantizer = new ProductQuantizer(this.optimizedConfig.productQuantization.numSubvectors, this.optimizedConfig.productQuantization.numCentroids);
|
|
}
|
|
// Reset memory usage (thread-safe)
|
|
const currentMemoryUsage = await this.getMemoryUsageAsync();
|
|
await this.updateMemoryUsage(-currentMemoryUsage.memoryUsage, -currentMemoryUsage.vectorCount);
|
|
// Clear the in-memory index
|
|
super.clear();
|
|
}
|
|
/**
|
|
* Initialize product quantizer with existing vectors
|
|
*/
|
|
initializeProductQuantizer() {
|
|
if (!this.productQuantizer) {
|
|
return;
|
|
}
|
|
// Get all vectors from the index
|
|
const nouns = super.getNouns();
|
|
const vectors = [];
|
|
// Extract vectors
|
|
for (const [_, noun] of nouns) {
|
|
vectors.push(noun.vector);
|
|
}
|
|
// Train the product quantizer
|
|
if (vectors.length > 0) {
|
|
this.productQuantizer.train(vectors);
|
|
// Quantize all existing vectors
|
|
for (const [id, noun] of nouns) {
|
|
const codes = this.productQuantizer.quantize(noun.vector);
|
|
this.quantizedVectors.set(id, codes);
|
|
}
|
|
console.log(`Initialized product quantizer with ${vectors.length} vectors`);
|
|
}
|
|
}
|
|
/**
|
|
* Get the product quantizer
|
|
* @returns Product quantizer or null if not enabled
|
|
*/
|
|
getProductQuantizer() {
|
|
return this.productQuantizer;
|
|
}
|
|
/**
|
|
* Get the optimized configuration
|
|
* @returns Optimized configuration
|
|
*/
|
|
getOptimizedConfig() {
|
|
return { ...this.optimizedConfig };
|
|
}
|
|
/**
|
|
* Get the estimated memory usage
|
|
* @returns Estimated memory usage in bytes
|
|
*/
|
|
getMemoryUsage() {
|
|
return this.memoryUsage;
|
|
}
|
|
/**
|
|
* Set the storage adapter
|
|
* @param storage Storage adapter
|
|
*/
|
|
setStorage(storage) {
|
|
this.storage = storage;
|
|
}
|
|
/**
|
|
* Get the storage adapter
|
|
* @returns Storage adapter or null if not set
|
|
*/
|
|
getStorage() {
|
|
return this.storage;
|
|
}
|
|
/**
|
|
* Set whether to use disk-based index
|
|
* @param useDiskBasedIndex Whether to use disk-based index
|
|
*/
|
|
setUseDiskBasedIndex(useDiskBasedIndex) {
|
|
this.useDiskBasedIndex = useDiskBasedIndex;
|
|
}
|
|
/**
|
|
* Get whether disk-based index is used
|
|
* @returns Whether disk-based index is used
|
|
*/
|
|
getUseDiskBasedIndex() {
|
|
return this.useDiskBasedIndex;
|
|
}
|
|
/**
|
|
* Set whether to use product quantization
|
|
* @param useProductQuantization Whether to use product quantization
|
|
*/
|
|
setUseProductQuantization(useProductQuantization) {
|
|
this.useProductQuantization = useProductQuantization;
|
|
}
|
|
/**
|
|
* Get whether product quantization is used
|
|
* @returns Whether product quantization is used
|
|
*/
|
|
getUseProductQuantization() {
|
|
return this.useProductQuantization;
|
|
}
|
|
}
|
|
//# sourceMappingURL=hnswIndexOptimized.js.map
|