**feat: integrate optimized index functionality into BrainyData**
### Changes: - Enhanced `BrainyData` class to support `HNSWIndexOptimized` alongside the standard HNSW index. - Added a new configuration parameter `hnswOptimized` for enabling optimized indexing. - Created utility methods for handling optimized index-specific functionalities such as memory usage tracking and disk-based index management. - Updated vector addition, retrieval, and management methods to accommodate both standard and optimized indices. - Introduced support for using pre-defined IDs during vector and relationship creation. - Standardized metadata handling for error resilience and embedding integrations. - Implemented a `getAllNouns` method to retrieve all indexed data. - Added validation for vector dimensions during relationship creation to ensure consistency. - Refactored method signatures to improve readability and alignment. ### Purpose: Improved the BrainyData API to support optimized indexing functionality powered by `HNSWIndexOptimized`, enhancing scalability for large datasets while maintaining backward compatibility. Elevated overall code structure and usability.
This commit is contained in:
parent
6f6f787bf1
commit
5703c4a69a
3 changed files with 1159 additions and 133 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -379,6 +379,34 @@ export class HNSWIndex {
|
|||
return this.distanceFunction
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the entry point ID
|
||||
*/
|
||||
public getEntryPointId(): string | null {
|
||||
return this.entryPointId
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the maximum level
|
||||
*/
|
||||
public getMaxLevel(): number {
|
||||
return this.maxLevel
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the dimension
|
||||
*/
|
||||
public getDimension(): number | null {
|
||||
return this.dimension
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the configuration
|
||||
*/
|
||||
public getConfig(): HNSWConfig {
|
||||
return { ...this.config }
|
||||
}
|
||||
|
||||
/**
|
||||
* Search within a specific layer
|
||||
* Returns a map of noun IDs to distances, sorted by distance
|
||||
|
|
|
|||
586
src/hnsw/hnswIndexOptimized.ts
Normal file
586
src/hnsw/hnswIndexOptimized.ts
Normal file
|
|
@ -0,0 +1,586 @@
|
|||
/**
|
||||
* 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 {
|
||||
DistanceFunction,
|
||||
HNSWConfig,
|
||||
HNSWNoun,
|
||||
Vector,
|
||||
VectorDocument
|
||||
} from '../coreTypes.js'
|
||||
import { HNSWIndex } from './hnswIndex.js'
|
||||
import { StorageAdapter } from '../coreTypes.js'
|
||||
|
||||
// Configuration for the optimized HNSW index
|
||||
export interface HNSWOptimizedConfig extends HNSWConfig {
|
||||
// Memory threshold in bytes - when exceeded, will use disk-based approach
|
||||
memoryThreshold?: number
|
||||
|
||||
// Product quantization settings
|
||||
productQuantization?: {
|
||||
// Whether to use product quantization
|
||||
enabled: boolean
|
||||
// Number of subvectors to split the vector into
|
||||
numSubvectors?: number
|
||||
// Number of centroids per subvector
|
||||
numCentroids?: number
|
||||
}
|
||||
|
||||
// Whether to use disk-based storage for the index
|
||||
useDiskBasedIndex?: boolean
|
||||
}
|
||||
|
||||
// Default configuration for the optimized HNSW index
|
||||
const DEFAULT_OPTIMIZED_CONFIG: HNSWOptimizedConfig = {
|
||||
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 {
|
||||
private numSubvectors: number
|
||||
private numCentroids: number
|
||||
private centroids: Vector[][] = []
|
||||
private subvectorSize: number = 0
|
||||
private initialized: boolean = false
|
||||
private dimension: number = 0
|
||||
|
||||
constructor(numSubvectors: number = 16, numCentroids: number = 256) {
|
||||
this.numSubvectors = numSubvectors
|
||||
this.numCentroids = numCentroids
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize the product quantizer with training data
|
||||
* @param vectors Training vectors to use for learning centroids
|
||||
*/
|
||||
public train(vectors: Vector[]): void {
|
||||
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: Vector[] = 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
|
||||
*/
|
||||
public quantize(vector: Vector): number[] {
|
||||
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: number[] = []
|
||||
|
||||
// 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
|
||||
*/
|
||||
public reconstruct(codes: number[]): Vector {
|
||||
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: Vector = []
|
||||
|
||||
// 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
|
||||
*/
|
||||
private euclideanDistanceSquared(a: Vector, b: Vector): number {
|
||||
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
|
||||
*/
|
||||
private kMeansPlusPlus(vectors: Vector[], k: number): Vector[] {
|
||||
if (vectors.length < k) {
|
||||
// If we have fewer vectors than centroids, use the vectors as centroids
|
||||
return [...vectors]
|
||||
}
|
||||
|
||||
const centroids: Vector[] = []
|
||||
|
||||
// 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: number[] = 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
|
||||
*/
|
||||
public getCentroids(): Vector[][] {
|
||||
return this.centroids
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the centroids for each subvector
|
||||
* @param centroids Array of centroid arrays
|
||||
*/
|
||||
public setCentroids(centroids: Vector[][]): void {
|
||||
this.centroids = centroids
|
||||
this.numSubvectors = centroids.length
|
||||
this.numCentroids = centroids[0].length
|
||||
this.initialized = true
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the dimension of the vectors
|
||||
* @returns Dimension
|
||||
*/
|
||||
public getDimension(): number {
|
||||
return this.dimension
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the dimension of the vectors
|
||||
* @param dimension Dimension
|
||||
*/
|
||||
public setDimension(dimension: number): void {
|
||||
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 {
|
||||
private optimizedConfig: HNSWOptimizedConfig
|
||||
private productQuantizer: ProductQuantizer | null = null
|
||||
private storage: StorageAdapter | null = null
|
||||
private useDiskBasedIndex: boolean = false
|
||||
private useProductQuantization: boolean = false
|
||||
private quantizedVectors: Map<string, number[]> = new Map()
|
||||
private memoryUsage: number = 0
|
||||
private vectorCount: number = 0
|
||||
|
||||
constructor(
|
||||
config: Partial<HNSWOptimizedConfig> = {},
|
||||
distanceFunction: DistanceFunction,
|
||||
storage: StorageAdapter | null = null
|
||||
) {
|
||||
// Initialize base HNSW index with standard config
|
||||
super(config, distanceFunction)
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a vector to the index
|
||||
* Uses product quantization if enabled and memory threshold is exceeded
|
||||
*/
|
||||
public override addItem(item: VectorDocument): string {
|
||||
// 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
|
||||
this.memoryUsage += totalMemory
|
||||
this.vectorCount++
|
||||
|
||||
// Check if we should switch to product quantization
|
||||
if (
|
||||
this.useProductQuantization &&
|
||||
this.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 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: HNSWNoun = {
|
||||
id,
|
||||
vector,
|
||||
connections: new Map()
|
||||
}
|
||||
|
||||
// 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 super.addItem(item)
|
||||
}
|
||||
|
||||
/**
|
||||
* Search for nearest neighbors
|
||||
* Uses product quantization if enabled
|
||||
*/
|
||||
public override search(
|
||||
queryVector: Vector,
|
||||
k: number = 10
|
||||
): Array<[string, number]> {
|
||||
// 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 super.search(reconstructedVector, k)
|
||||
}
|
||||
|
||||
// Otherwise, use the standard search
|
||||
return super.search(queryVector, k)
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove an item from the index
|
||||
*/
|
||||
public override removeItem(id: string): boolean {
|
||||
// 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
|
||||
if (this.vectorCount > 0) {
|
||||
this.memoryUsage = Math.max(
|
||||
0,
|
||||
this.memoryUsage - this.memoryUsage / this.vectorCount
|
||||
)
|
||||
this.vectorCount--
|
||||
}
|
||||
|
||||
// Remove the item from the in-memory index
|
||||
return super.removeItem(id)
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear the index
|
||||
*/
|
||||
public override clear(): void {
|
||||
// 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
|
||||
this.memoryUsage = 0
|
||||
this.vectorCount = 0
|
||||
|
||||
// Clear the in-memory index
|
||||
super.clear()
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize product quantizer with existing vectors
|
||||
*/
|
||||
private initializeProductQuantizer(): void {
|
||||
if (!this.productQuantizer) {
|
||||
return
|
||||
}
|
||||
|
||||
// Get all vectors from the index
|
||||
const nouns = super.getNouns()
|
||||
const vectors: Vector[] = []
|
||||
|
||||
// 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
|
||||
*/
|
||||
public getProductQuantizer(): ProductQuantizer | null {
|
||||
return this.productQuantizer
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the optimized configuration
|
||||
* @returns Optimized configuration
|
||||
*/
|
||||
public getOptimizedConfig(): HNSWOptimizedConfig {
|
||||
return { ...this.optimizedConfig }
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the estimated memory usage
|
||||
* @returns Estimated memory usage in bytes
|
||||
*/
|
||||
public getMemoryUsage(): number {
|
||||
return this.memoryUsage
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the storage adapter
|
||||
* @param storage Storage adapter
|
||||
*/
|
||||
public setStorage(storage: StorageAdapter): void {
|
||||
this.storage = storage
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the storage adapter
|
||||
* @returns Storage adapter or null if not set
|
||||
*/
|
||||
public getStorage(): StorageAdapter | null {
|
||||
return this.storage
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to use disk-based index
|
||||
* @param useDiskBasedIndex Whether to use disk-based index
|
||||
*/
|
||||
public setUseDiskBasedIndex(useDiskBasedIndex: boolean): void {
|
||||
this.useDiskBasedIndex = useDiskBasedIndex
|
||||
}
|
||||
|
||||
/**
|
||||
* Get whether disk-based index is used
|
||||
* @returns Whether disk-based index is used
|
||||
*/
|
||||
public getUseDiskBasedIndex(): boolean {
|
||||
return this.useDiskBasedIndex
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to use product quantization
|
||||
* @param useProductQuantization Whether to use product quantization
|
||||
*/
|
||||
public setUseProductQuantization(useProductQuantization: boolean): void {
|
||||
this.useProductQuantization = useProductQuantization
|
||||
}
|
||||
|
||||
/**
|
||||
* Get whether product quantization is used
|
||||
* @returns Whether product quantization is used
|
||||
*/
|
||||
public getUseProductQuantization(): boolean {
|
||||
return this.useProductQuantization
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue