Initial commit
This commit is contained in:
commit
5a8a6c1ba3
81 changed files with 39269 additions and 0 deletions
704
src/hnsw/hnswIndex.ts
Normal file
704
src/hnsw/hnswIndex.ts
Normal file
|
|
@ -0,0 +1,704 @@
|
|||
/**
|
||||
* HNSW (Hierarchical Navigable Small World) Index implementation
|
||||
* Based on the paper: "Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs"
|
||||
*/
|
||||
|
||||
import { DistanceFunction, HNSWConfig, HNSWNoun, Vector, VectorDocument } from '../coreTypes.js'
|
||||
import { euclideanDistance } from '../utils/index.js'
|
||||
import { executeInThread } from '../utils/workerUtils.js'
|
||||
|
||||
// Default HNSW parameters
|
||||
const DEFAULT_CONFIG: HNSWConfig = {
|
||||
M: 16, // Max number of connections per noun
|
||||
efConstruction: 200, // Size of a dynamic candidate list during construction
|
||||
efSearch: 50, // Size of a dynamic candidate list during search
|
||||
ml: 16 // Max level
|
||||
}
|
||||
|
||||
export class HNSWIndex {
|
||||
private nouns: Map<string, HNSWNoun> = new Map()
|
||||
private entryPointId: string | null = null
|
||||
private maxLevel = 0
|
||||
private config: HNSWConfig
|
||||
private distanceFunction: DistanceFunction
|
||||
private dimension: number | null = null
|
||||
private useParallelization: boolean = true // Whether to use parallelization for performance-critical operations
|
||||
|
||||
constructor(
|
||||
config: Partial<HNSWConfig> = {},
|
||||
distanceFunction: DistanceFunction = euclideanDistance,
|
||||
options: { useParallelization?: boolean } = {}
|
||||
) {
|
||||
this.config = { ...DEFAULT_CONFIG, ...config }
|
||||
this.distanceFunction = distanceFunction
|
||||
this.useParallelization = options.useParallelization !== undefined ? options.useParallelization : true
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to use parallelization for performance-critical operations
|
||||
*/
|
||||
public setUseParallelization(useParallelization: boolean): void {
|
||||
this.useParallelization = useParallelization
|
||||
}
|
||||
|
||||
/**
|
||||
* Get whether parallelization is enabled
|
||||
*/
|
||||
public getUseParallelization(): boolean {
|
||||
return this.useParallelization
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculate distances between a query vector and multiple vectors in parallel
|
||||
* This is used to optimize performance for search operations
|
||||
* @param queryVector The query vector
|
||||
* @param vectors Array of vectors to compare against
|
||||
* @returns Array of distances
|
||||
*/
|
||||
private async calculateDistancesInParallel(
|
||||
queryVector: Vector,
|
||||
vectors: Array<{ id: string; vector: Vector }>
|
||||
): Promise<Array<{ id: string; distance: number }>> {
|
||||
// If parallelization is disabled or there are very few vectors, use sequential processing
|
||||
if (!this.useParallelization || vectors.length < 10) {
|
||||
return vectors.map(item => ({
|
||||
id: item.id,
|
||||
distance: this.distanceFunction(queryVector, item.vector)
|
||||
}))
|
||||
}
|
||||
|
||||
// Function to be executed in a worker thread
|
||||
const distanceCalculator = (
|
||||
args: {
|
||||
queryVector: Vector,
|
||||
vectors: Array<{ id: string; vector: Vector }>,
|
||||
distanceFnString: string
|
||||
}
|
||||
) => {
|
||||
const { queryVector, vectors, distanceFnString } = args;
|
||||
|
||||
// Recreate the distance function from its string representation
|
||||
const distanceFunction = new Function('return ' + distanceFnString)() as DistanceFunction
|
||||
|
||||
// Calculate distances for all items
|
||||
return vectors.map(item => ({
|
||||
id: item.id,
|
||||
distance: distanceFunction(queryVector, item.vector)
|
||||
}))
|
||||
}
|
||||
|
||||
try {
|
||||
// Convert the distance function to a string for serialization
|
||||
const distanceFnString = this.distanceFunction.toString()
|
||||
|
||||
// Execute the distance calculation in a separate thread
|
||||
return await executeInThread<Array<{ id: string; distance: number }>>(
|
||||
distanceCalculator.toString(),
|
||||
{ queryVector, vectors, distanceFnString }
|
||||
)
|
||||
} catch (error) {
|
||||
console.error('Error in parallel distance calculation, falling back to sequential:', error)
|
||||
// Fall back to sequential processing if parallel execution fails
|
||||
return vectors.map(item => ({
|
||||
id: item.id,
|
||||
distance: this.distanceFunction(queryVector, item.vector)
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a vector to the index
|
||||
*/
|
||||
public async addItem(item: VectorDocument): Promise<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')
|
||||
}
|
||||
|
||||
// Set dimension on first insert
|
||||
if (this.dimension === null) {
|
||||
this.dimension = vector.length
|
||||
} else if (vector.length !== this.dimension) {
|
||||
throw new Error(
|
||||
`Vector dimension mismatch: expected ${this.dimension}, got ${vector.length}`
|
||||
)
|
||||
}
|
||||
|
||||
// Generate random level for this noun
|
||||
const nounLevel = this.getRandomLevel()
|
||||
|
||||
// Create new noun
|
||||
const noun: HNSWNoun = {
|
||||
id,
|
||||
vector,
|
||||
connections: new Map()
|
||||
}
|
||||
|
||||
// Initialize empty connection sets for each level
|
||||
for (let level = 0; level <= nounLevel; level++) {
|
||||
noun.connections.set(level, new Set<string>())
|
||||
}
|
||||
|
||||
// If this is the first noun, make it the entry point
|
||||
if (this.nouns.size === 0) {
|
||||
this.entryPointId = id
|
||||
this.maxLevel = nounLevel
|
||||
this.nouns.set(id, noun)
|
||||
return id
|
||||
}
|
||||
|
||||
// Find entry point
|
||||
if (!this.entryPointId) {
|
||||
console.error('Entry point ID is null')
|
||||
// If there's no entry point, this is the first noun, so we should have returned earlier
|
||||
// This is a safety check
|
||||
this.entryPointId = id
|
||||
this.maxLevel = nounLevel
|
||||
this.nouns.set(id, noun)
|
||||
return id
|
||||
}
|
||||
|
||||
const entryPoint = this.nouns.get(this.entryPointId)
|
||||
if (!entryPoint) {
|
||||
console.error(`Entry point with ID ${this.entryPointId} not found`)
|
||||
// If the entry point doesn't exist, treat this as the first noun
|
||||
this.entryPointId = id
|
||||
this.maxLevel = nounLevel
|
||||
this.nouns.set(id, noun)
|
||||
return id
|
||||
}
|
||||
|
||||
let currObj = entryPoint
|
||||
let currDist = this.distanceFunction(vector, entryPoint.vector)
|
||||
|
||||
// Traverse the graph from top to bottom to find the closest noun
|
||||
for (let level = this.maxLevel; level > nounLevel; level--) {
|
||||
let changed = true
|
||||
while (changed) {
|
||||
changed = false
|
||||
|
||||
// Check all neighbors at current level
|
||||
const connections = currObj.connections.get(level) || new Set<string>()
|
||||
|
||||
for (const neighborId of connections) {
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) {
|
||||
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
||||
continue
|
||||
}
|
||||
const distToNeighbor = this.distanceFunction(vector, neighbor.vector)
|
||||
|
||||
if (distToNeighbor < currDist) {
|
||||
currDist = distToNeighbor
|
||||
currObj = neighbor
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For each level from nounLevel down to 0
|
||||
for (let level = Math.min(nounLevel, this.maxLevel); level >= 0; level--) {
|
||||
// Find ef nearest elements using greedy search
|
||||
const nearestNouns = await this.searchLayer(
|
||||
vector,
|
||||
currObj,
|
||||
this.config.efConstruction,
|
||||
level
|
||||
)
|
||||
|
||||
// Select M nearest neighbors
|
||||
const neighbors = this.selectNeighbors(
|
||||
vector,
|
||||
nearestNouns,
|
||||
this.config.M
|
||||
)
|
||||
|
||||
// Add bidirectional connections
|
||||
for (const [neighborId, _] of neighbors) {
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) {
|
||||
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
||||
continue
|
||||
}
|
||||
|
||||
noun.connections.get(level)!.add(neighborId)
|
||||
|
||||
// Add reverse connection
|
||||
if (!neighbor.connections.has(level)) {
|
||||
neighbor.connections.set(level, new Set<string>())
|
||||
}
|
||||
neighbor.connections.get(level)!.add(id)
|
||||
|
||||
// Ensure neighbor doesn't have too many connections
|
||||
if (neighbor.connections.get(level)!.size > this.config.M) {
|
||||
this.pruneConnections(neighbor, level)
|
||||
}
|
||||
}
|
||||
|
||||
// Update entry point for the next level
|
||||
if (nearestNouns.size > 0) {
|
||||
const [nearestId, nearestDist] = [...nearestNouns][0]
|
||||
if (nearestDist < currDist) {
|
||||
currDist = nearestDist
|
||||
const nearestNoun = this.nouns.get(nearestId)
|
||||
if (!nearestNoun) {
|
||||
console.error(`Nearest noun with ID ${nearestId} not found in addItem`)
|
||||
// Keep the current object as is
|
||||
} else {
|
||||
currObj = nearestNoun
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update max level and entry point if needed
|
||||
if (nounLevel > this.maxLevel) {
|
||||
this.maxLevel = nounLevel
|
||||
this.entryPointId = id
|
||||
}
|
||||
|
||||
// Add noun to the index
|
||||
this.nouns.set(id, noun)
|
||||
return id
|
||||
}
|
||||
|
||||
/**
|
||||
* Search for nearest neighbors
|
||||
*/
|
||||
public async search(queryVector: Vector, k: number = 10): Promise<Array<[string, number]>> {
|
||||
if (this.nouns.size === 0) {
|
||||
return []
|
||||
}
|
||||
|
||||
// Check if query vector is defined
|
||||
if (!queryVector) {
|
||||
throw new Error('Query vector is undefined or null')
|
||||
}
|
||||
|
||||
if (this.dimension !== null && queryVector.length !== this.dimension) {
|
||||
throw new Error(
|
||||
`Query vector dimension mismatch: expected ${this.dimension}, got ${queryVector.length}`
|
||||
)
|
||||
}
|
||||
|
||||
// Start from the entry point
|
||||
if (!this.entryPointId) {
|
||||
console.error('Entry point ID is null')
|
||||
return []
|
||||
}
|
||||
|
||||
const entryPoint = this.nouns.get(this.entryPointId)
|
||||
if (!entryPoint) {
|
||||
console.error(`Entry point with ID ${this.entryPointId} not found`)
|
||||
return []
|
||||
}
|
||||
|
||||
let currObj = entryPoint
|
||||
let currDist = this.distanceFunction(queryVector, currObj.vector)
|
||||
|
||||
// Traverse the graph from top to bottom to find the closest noun
|
||||
for (let level = this.maxLevel; level > 0; level--) {
|
||||
let changed = true
|
||||
while (changed) {
|
||||
changed = false
|
||||
|
||||
// Check all neighbors at current level
|
||||
const connections = currObj.connections.get(level) || new Set<string>()
|
||||
|
||||
// If we have enough connections, use parallel distance calculation
|
||||
if (this.useParallelization && connections.size >= 10) {
|
||||
// Prepare vectors for parallel calculation
|
||||
const vectors: Array<{ id: string; vector: Vector }> = []
|
||||
for (const neighborId of connections) {
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) continue
|
||||
vectors.push({ id: neighborId, vector: neighbor.vector })
|
||||
}
|
||||
|
||||
// Calculate distances in parallel
|
||||
const distances = await this.calculateDistancesInParallel(queryVector, vectors)
|
||||
|
||||
// Find the closest neighbor
|
||||
for (const { id, distance } of distances) {
|
||||
if (distance < currDist) {
|
||||
currDist = distance
|
||||
const neighbor = this.nouns.get(id)
|
||||
if (neighbor) {
|
||||
currObj = neighbor
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use sequential processing for small number of connections
|
||||
for (const neighborId of connections) {
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) {
|
||||
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
||||
continue
|
||||
}
|
||||
const distToNeighbor = this.distanceFunction(
|
||||
queryVector,
|
||||
neighbor.vector
|
||||
)
|
||||
|
||||
if (distToNeighbor < currDist) {
|
||||
currDist = distToNeighbor
|
||||
currObj = neighbor
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Search at level 0 with ef = k
|
||||
const nearestNouns = await this.searchLayer(
|
||||
queryVector,
|
||||
currObj,
|
||||
Math.max(this.config.efSearch, k),
|
||||
0
|
||||
)
|
||||
|
||||
// Convert to array and sort by distance
|
||||
return [...nearestNouns].slice(0, k)
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove an item from the index
|
||||
*/
|
||||
public removeItem(id: string): boolean {
|
||||
if (!this.nouns.has(id)) {
|
||||
return false
|
||||
}
|
||||
|
||||
const noun = this.nouns.get(id)!
|
||||
|
||||
// Remove connections to this noun from all neighbors
|
||||
for (const [level, connections] of noun.connections.entries()) {
|
||||
for (const neighborId of connections) {
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) {
|
||||
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
||||
continue
|
||||
}
|
||||
if (neighbor.connections.has(level)) {
|
||||
neighbor.connections.get(level)!.delete(id)
|
||||
|
||||
// Prune connections after removing this noun to ensure consistency
|
||||
this.pruneConnections(neighbor, level)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Also check all other nouns for references to this noun and remove them
|
||||
for (const [nounId, otherNoun] of this.nouns.entries()) {
|
||||
if (nounId === id) continue // Skip the noun being removed
|
||||
|
||||
for (const [level, connections] of otherNoun.connections.entries()) {
|
||||
if (connections.has(id)) {
|
||||
connections.delete(id)
|
||||
|
||||
// Prune connections after removing this reference
|
||||
this.pruneConnections(otherNoun, level)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove the noun
|
||||
this.nouns.delete(id)
|
||||
|
||||
// If we removed the entry point, find a new one
|
||||
if (this.entryPointId === id) {
|
||||
if (this.nouns.size === 0) {
|
||||
this.entryPointId = null
|
||||
this.maxLevel = 0
|
||||
} else {
|
||||
// Find the noun with the highest level
|
||||
let maxLevel = 0
|
||||
let newEntryPointId = null
|
||||
|
||||
for (const [nounId, noun] of this.nouns.entries()) {
|
||||
if (noun.connections.size === 0) continue // Skip nouns with no connections
|
||||
|
||||
const nounLevel = Math.max(...noun.connections.keys())
|
||||
if (nounLevel >= maxLevel) {
|
||||
maxLevel = nounLevel
|
||||
newEntryPointId = nounId
|
||||
}
|
||||
}
|
||||
|
||||
this.entryPointId = newEntryPointId
|
||||
this.maxLevel = maxLevel
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all nouns in the index
|
||||
*/
|
||||
public getNouns(): Map<string, HNSWNoun> {
|
||||
return new Map(this.nouns)
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Clear the index
|
||||
*/
|
||||
public clear(): void {
|
||||
this.nouns.clear()
|
||||
this.entryPointId = null
|
||||
this.maxLevel = 0
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the size of the index
|
||||
*/
|
||||
public size(): number {
|
||||
return this.nouns.size
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the distance function used by the index
|
||||
*/
|
||||
public getDistanceFunction(): DistanceFunction {
|
||||
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
|
||||
*/
|
||||
private async searchLayer(
|
||||
queryVector: Vector,
|
||||
entryPoint: HNSWNoun,
|
||||
ef: number,
|
||||
level: number
|
||||
): Promise<Map<string, number>> {
|
||||
// Set of visited nouns
|
||||
const visited = new Set<string>([entryPoint.id])
|
||||
|
||||
// Priority queue of candidates (closest first)
|
||||
const candidates = new Map<string, number>()
|
||||
candidates.set(
|
||||
entryPoint.id,
|
||||
this.distanceFunction(queryVector, entryPoint.vector)
|
||||
)
|
||||
|
||||
// Priority queue of nearest neighbors found so far (closest first)
|
||||
const nearest = new Map<string, number>()
|
||||
nearest.set(
|
||||
entryPoint.id,
|
||||
this.distanceFunction(queryVector, entryPoint.vector)
|
||||
)
|
||||
|
||||
// While there are candidates to explore
|
||||
while (candidates.size > 0) {
|
||||
// Get closest candidate
|
||||
const [closestId, closestDist] = [...candidates][0]
|
||||
candidates.delete(closestId)
|
||||
|
||||
// If this candidate is farther than the farthest in our result set, we're done
|
||||
const farthestInNearest = [...nearest][nearest.size - 1]
|
||||
if (nearest.size >= ef && closestDist > farthestInNearest[1]) {
|
||||
break
|
||||
}
|
||||
|
||||
// Explore neighbors of the closest candidate
|
||||
const noun = this.nouns.get(closestId)
|
||||
if (!noun) {
|
||||
console.error(`Noun with ID ${closestId} not found in searchLayer`)
|
||||
continue
|
||||
}
|
||||
const connections = noun.connections.get(level) || new Set<string>()
|
||||
|
||||
// If we have enough connections and parallelization is enabled, use parallel distance calculation
|
||||
if (this.useParallelization && connections.size >= 10) {
|
||||
// Collect unvisited neighbors
|
||||
const unvisitedNeighbors: Array<{ id: string; vector: Vector }> = []
|
||||
for (const neighborId of connections) {
|
||||
if (!visited.has(neighborId)) {
|
||||
visited.add(neighborId)
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) continue
|
||||
unvisitedNeighbors.push({ id: neighborId, vector: neighbor.vector })
|
||||
}
|
||||
}
|
||||
|
||||
if (unvisitedNeighbors.length > 0) {
|
||||
// Calculate distances in parallel
|
||||
const distances = await this.calculateDistancesInParallel(queryVector, unvisitedNeighbors)
|
||||
|
||||
// Process the results
|
||||
for (const { id, distance } of distances) {
|
||||
// If we haven't found ef nearest neighbors yet, or this neighbor is closer than the farthest one we've found
|
||||
if (nearest.size < ef || distance < farthestInNearest[1]) {
|
||||
candidates.set(id, distance)
|
||||
nearest.set(id, distance)
|
||||
|
||||
// If we have more than ef neighbors, remove the farthest one
|
||||
if (nearest.size > ef) {
|
||||
const sortedNearest = [...nearest].sort((a, b) => a[1] - b[1])
|
||||
nearest.clear()
|
||||
for (let i = 0; i < ef; i++) {
|
||||
nearest.set(sortedNearest[i][0], sortedNearest[i][1])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use sequential processing for small number of connections
|
||||
for (const neighborId of connections) {
|
||||
if (!visited.has(neighborId)) {
|
||||
visited.add(neighborId)
|
||||
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) {
|
||||
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
||||
continue
|
||||
}
|
||||
const distToNeighbor = this.distanceFunction(
|
||||
queryVector,
|
||||
neighbor.vector
|
||||
)
|
||||
|
||||
// If we haven't found ef nearest neighbors yet, or this neighbor is closer than the farthest one we've found
|
||||
if (nearest.size < ef || distToNeighbor < farthestInNearest[1]) {
|
||||
candidates.set(neighborId, distToNeighbor)
|
||||
nearest.set(neighborId, distToNeighbor)
|
||||
|
||||
// If we have more than ef neighbors, remove the farthest one
|
||||
if (nearest.size > ef) {
|
||||
const sortedNearest = [...nearest].sort((a, b) => a[1] - b[1])
|
||||
nearest.clear()
|
||||
for (let i = 0; i < ef; i++) {
|
||||
nearest.set(sortedNearest[i][0], sortedNearest[i][1])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort nearest by distance
|
||||
return new Map([...nearest].sort((a, b) => a[1] - b[1]))
|
||||
}
|
||||
|
||||
/**
|
||||
* Select M nearest neighbors from the candidate set
|
||||
*/
|
||||
private selectNeighbors(
|
||||
queryVector: Vector,
|
||||
candidates: Map<string, number>,
|
||||
M: number
|
||||
): Map<string, number> {
|
||||
if (candidates.size <= M) {
|
||||
return candidates
|
||||
}
|
||||
|
||||
// Simple heuristic: just take the M closest
|
||||
const sortedCandidates = [...candidates].sort((a, b) => a[1] - b[1])
|
||||
const result = new Map<string, number>()
|
||||
|
||||
for (let i = 0; i < Math.min(M, sortedCandidates.length); i++) {
|
||||
result.set(sortedCandidates[i][0], sortedCandidates[i][1])
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
* Ensure a noun doesn't have too many connections at a given level
|
||||
*/
|
||||
private pruneConnections(noun: HNSWNoun, level: number): void {
|
||||
const connections = noun.connections.get(level)!
|
||||
if (connections.size <= this.config.M) {
|
||||
return
|
||||
}
|
||||
|
||||
// Calculate distances to all neighbors
|
||||
const distances = new Map<string, number>()
|
||||
const validNeighborIds = new Set<string>()
|
||||
|
||||
for (const neighborId of connections) {
|
||||
const neighbor = this.nouns.get(neighborId)
|
||||
if (!neighbor) {
|
||||
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
||||
continue
|
||||
}
|
||||
|
||||
// Only add valid neighbors to the distances map
|
||||
distances.set(
|
||||
neighborId,
|
||||
this.distanceFunction(noun.vector, neighbor.vector)
|
||||
)
|
||||
validNeighborIds.add(neighborId)
|
||||
}
|
||||
|
||||
// Only proceed if we have valid neighbors
|
||||
if (distances.size === 0) {
|
||||
// If no valid neighbors, clear connections at this level
|
||||
noun.connections.set(level, new Set())
|
||||
return
|
||||
}
|
||||
|
||||
// Select M closest neighbors from valid ones
|
||||
const selectedNeighbors = this.selectNeighbors(
|
||||
noun.vector,
|
||||
distances,
|
||||
this.config.M
|
||||
)
|
||||
|
||||
// Update connections with only valid neighbors
|
||||
noun.connections.set(level, new Set(selectedNeighbors.keys()))
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a random level for a new noun
|
||||
* Uses the same distribution as in the original HNSW paper
|
||||
*/
|
||||
private getRandomLevel(): number {
|
||||
const r = Math.random()
|
||||
return Math.floor(-Math.log(r) * (1.0 / Math.log(this.config.M)))
|
||||
}
|
||||
}
|
||||
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 async addItem(item: VectorDocument): Promise<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 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: 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 await super.addItem(item)
|
||||
}
|
||||
|
||||
/**
|
||||
* Search for nearest neighbors
|
||||
* Uses product quantization if enabled
|
||||
*/
|
||||
public override async search(
|
||||
queryVector: Vector,
|
||||
k: number = 10
|
||||
): Promise<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 await super.search(reconstructedVector, k)
|
||||
}
|
||||
|
||||
// Otherwise, use the standard search
|
||||
return await 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