New memory system
This commit is contained in:
334
src/kd-tree.ts
Normal file
334
src/kd-tree.ts
Normal file
@@ -0,0 +1,334 @@
|
||||
export type DistanceMetric = "euclidean" | "cosine";
|
||||
|
||||
export interface KDPoint<T = unknown> {
|
||||
vector: number[];
|
||||
payload: T;
|
||||
}
|
||||
|
||||
export interface KNNResult<T = unknown> {
|
||||
point: KDPoint<T>;
|
||||
distance: number;
|
||||
}
|
||||
|
||||
interface KDNode<T> {
|
||||
point: KDPoint<T>;
|
||||
axis: number;
|
||||
left: KDNode<T> | null;
|
||||
right: KDNode<T> | null;
|
||||
}
|
||||
|
||||
// ─── Distance helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
function euclidean(a: number[], b: number[]): number {
|
||||
let sum = 0;
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
const d = a[i] - b[i];
|
||||
sum += d * d;
|
||||
}
|
||||
return Math.sqrt(sum);
|
||||
}
|
||||
|
||||
function cosine(a: number[], b: number[]): number {
|
||||
let dot = 0, normA = 0, 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);
|
||||
return denom === 0 ? 1 : 1 - dot / denom; // distance = 1 - similarity
|
||||
}
|
||||
|
||||
/**
|
||||
* Keeps the k closest candidates in memory, evicts the furthest when full
|
||||
*/
|
||||
class BoundedMaxHeap<T> {
|
||||
private heap: KNNResult<T>[] = [];
|
||||
|
||||
constructor(private readonly k: number) {}
|
||||
|
||||
get size(): number { return this.heap.length; }
|
||||
|
||||
get worstDistance(): number {
|
||||
return this.heap.length < this.k ? Infinity : this.heap[0].distance;
|
||||
}
|
||||
|
||||
push(item: KNNResult<T>): void {
|
||||
if (this.heap.length < this.k) {
|
||||
this.heap.push(item);
|
||||
this.bubbleUp(this.heap.length - 1);
|
||||
} else if (item.distance < this.heap[0].distance) {
|
||||
this.heap[0] = item;
|
||||
this.sinkDown(0);
|
||||
}
|
||||
}
|
||||
|
||||
toSortedArray(): KNNResult<T>[] {
|
||||
return [...this.heap].sort((a, b) => a.distance - b.distance);
|
||||
}
|
||||
|
||||
private bubbleUp(i: number): void {
|
||||
while (i > 0) {
|
||||
const parent = (i - 1) >> 1;
|
||||
if (this.heap[parent].distance >= this.heap[i].distance) break;
|
||||
[this.heap[parent], this.heap[i]] = [this.heap[i], this.heap[parent]];
|
||||
i = parent;
|
||||
}
|
||||
}
|
||||
|
||||
private sinkDown(i: number): void {
|
||||
const n = this.heap.length;
|
||||
while (true) {
|
||||
let largest = i;
|
||||
const l = 2 * i + 1, r = 2 * i + 2;
|
||||
if (l < n && this.heap[l].distance > this.heap[largest].distance) largest = l;
|
||||
if (r < n && this.heap[r].distance > this.heap[largest].distance) largest = r;
|
||||
if (largest === i) break;
|
||||
[this.heap[largest], this.heap[i]] = [this.heap[i], this.heap[largest]];
|
||||
i = largest;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* K-D Tree for efficient nearest-neighbor search over high-dimensional vectors / embeddings.
|
||||
*
|
||||
* Supports:
|
||||
* - Insertion of labeled points
|
||||
* - k-nearest-neighbor (KNN) search
|
||||
* - Radius search (all points within a given distance)
|
||||
* - Euclidean and cosine distance metrics
|
||||
* - Bulk construction (balanced tree) for best query performance
|
||||
*/
|
||||
export class KDTree<T = unknown> {
|
||||
private root: KDNode<T> | null = null;
|
||||
private _size = 0;
|
||||
private readonly dims: number;
|
||||
private readonly distanceFn: (a: number[], b: number[]) => number;
|
||||
|
||||
/**
|
||||
* @param dims Dimensionality of all vectors (must be consistent).
|
||||
* @param metric Distance metric to use. Default: "euclidean".
|
||||
* @param points Optional initial set of points. Builds a balanced tree
|
||||
* in O(n log² n) — prefer this over inserting one-by-one
|
||||
* when you have a large corpus.
|
||||
*/
|
||||
constructor(
|
||||
dims: number,
|
||||
metric: DistanceMetric = "euclidean",
|
||||
points?: KDPoint<T>[]
|
||||
) {
|
||||
this.dims = dims;
|
||||
this.distanceFn = metric === "cosine" ? cosine : euclidean;
|
||||
|
||||
if (points && points.length > 0) {
|
||||
this.validateAll(points);
|
||||
this.root = this.buildBalanced([...points], 0);
|
||||
this._size = points.length;
|
||||
}
|
||||
}
|
||||
|
||||
/** Total number of points stored in the tree. */
|
||||
get size(): number { return this._size; }
|
||||
|
||||
// ── Insertion ──────────────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Insert a single point. O(log n) average, O(n) worst case on skewed data.
|
||||
* For bulk loading prefer passing points to the constructor.
|
||||
*/
|
||||
insert(point: KDPoint<T>): void {
|
||||
this.validate(point);
|
||||
this.root = this.insertNode(this.root, point, 0);
|
||||
this._size++;
|
||||
}
|
||||
|
||||
// ── KNN search ─────────────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Find the k nearest neighbors to `query`.
|
||||
* Returns results sorted by distance ascending.
|
||||
*/
|
||||
knn(query: number[], k: number): KNNResult<T>[] {
|
||||
if (k <= 0) throw new RangeError("k must be a positive integer");
|
||||
this.validateVector(query);
|
||||
|
||||
const heap = new BoundedMaxHeap<T>(k);
|
||||
this.searchKNN(this.root, query, k, heap, 0);
|
||||
return heap.toSortedArray();
|
||||
}
|
||||
|
||||
/**
|
||||
* Nearest single neighbor. Convenience wrapper around knn(query, 1).
|
||||
* Returns null if the tree is empty.
|
||||
*/
|
||||
nearest(query: number[]): KNNResult<T> | null {
|
||||
const results = this.knn(query, 1);
|
||||
return results[0] ?? null;
|
||||
}
|
||||
|
||||
// ── Radius search ──────────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* Return all points whose distance to `query` is ≤ `radius`,
|
||||
* sorted by distance ascending.
|
||||
*/
|
||||
radiusSearch(query: number[], radius: number): KNNResult<T>[] {
|
||||
if (radius < 0) throw new RangeError("radius must be non-negative");
|
||||
this.validateVector(query);
|
||||
|
||||
const results: KNNResult<T>[] = [];
|
||||
this.searchRadius(this.root, query, radius, results, 0);
|
||||
results.sort((a, b) => a.distance - b.distance);
|
||||
return results;
|
||||
}
|
||||
|
||||
// ── Conversion ─────────────────────────────────────────────────────────────
|
||||
|
||||
/** Collect all points in the tree (order not guaranteed). */
|
||||
toArray(): KDPoint<T>[] {
|
||||
const out: KDPoint<T>[] = [];
|
||||
this.collect(this.root, out);
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* Rebuild the tree from its current points as a balanced tree.
|
||||
* Useful after many individual insertions to restore O(log n) query time.
|
||||
*/
|
||||
rebalance(): void {
|
||||
const points = this.toArray();
|
||||
this.root = points.length ? this.buildBalanced(points, 0) : null;
|
||||
}
|
||||
|
||||
// ── Private: build ─────────────────────────────────────────────────────────
|
||||
|
||||
private buildBalanced(points: KDPoint<T>[], depth: number): KDNode<T> {
|
||||
const axis = depth % this.dims;
|
||||
points.sort((a, b) => a.vector[axis] - b.vector[axis]);
|
||||
|
||||
const mid = Math.floor(points.length / 2);
|
||||
return {
|
||||
point: points[mid],
|
||||
axis,
|
||||
left: points.slice(0, mid).length
|
||||
? this.buildBalanced(points.slice(0, mid), depth + 1)
|
||||
: null,
|
||||
right: points.slice(mid + 1).length
|
||||
? this.buildBalanced(points.slice(mid + 1), depth + 1)
|
||||
: null,
|
||||
};
|
||||
}
|
||||
|
||||
// ── Private: insert ────────────────────────────────────────────────────────
|
||||
|
||||
private insertNode(
|
||||
node: KDNode<T> | null,
|
||||
point: KDPoint<T>,
|
||||
depth: number
|
||||
): KDNode<T> {
|
||||
if (node === null) {
|
||||
return { point, axis: depth % this.dims, left: null, right: null };
|
||||
}
|
||||
const axis = depth % this.dims;
|
||||
if (point.vector[axis] < node.point.vector[axis]) {
|
||||
node.left = this.insertNode(node.left, point, depth + 1);
|
||||
} else {
|
||||
node.right = this.insertNode(node.right, point, depth + 1);
|
||||
}
|
||||
return node;
|
||||
}
|
||||
|
||||
// ── Private: KNN traversal ─────────────────────────────────────────────────
|
||||
|
||||
private searchKNN(
|
||||
node: KDNode<T> | null,
|
||||
query: number[],
|
||||
k: number,
|
||||
heap: BoundedMaxHeap<T>,
|
||||
depth: number
|
||||
): void {
|
||||
if (node === null) return;
|
||||
|
||||
const dist = this.distanceFn(query, node.point.vector);
|
||||
heap.push({ point: node.point, distance: dist });
|
||||
|
||||
const axis = node.axis;
|
||||
const diff = query[axis] - node.point.vector[axis];
|
||||
const [near, far] = diff <= 0
|
||||
? [node.left, node.right]
|
||||
: [node.right, node.left];
|
||||
|
||||
this.searchKNN(near, query, k, heap, depth + 1);
|
||||
|
||||
// Only explore the far side if it could contain a closer point.
|
||||
// For cosine distance we can't prune by axis gap alone, so always explore.
|
||||
const shouldExplore =
|
||||
this.distanceFn === cosine
|
||||
? true
|
||||
: Math.abs(diff) < heap.worstDistance;
|
||||
|
||||
if (shouldExplore) {
|
||||
this.searchKNN(far, query, k, heap, depth + 1);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Private: radius traversal ──────────────────────────────────────────────
|
||||
|
||||
private searchRadius(
|
||||
node: KDNode<T> | null,
|
||||
query: number[],
|
||||
radius: number,
|
||||
results: KNNResult<T>[],
|
||||
depth: number
|
||||
): void {
|
||||
if (node === null) return;
|
||||
|
||||
const dist = this.distanceFn(query, node.point.vector);
|
||||
if (dist <= radius) {
|
||||
results.push({ point: node.point, distance: dist });
|
||||
}
|
||||
|
||||
const axis = node.axis;
|
||||
const diff = query[axis] - node.point.vector[axis];
|
||||
const [near, far] = diff <= 0
|
||||
? [node.left, node.right]
|
||||
: [node.right, node.left];
|
||||
|
||||
this.searchRadius(near, query, radius, results, depth + 1);
|
||||
|
||||
const shouldExplore =
|
||||
this.distanceFn === cosine ? true : Math.abs(diff) <= radius;
|
||||
|
||||
if (shouldExplore) {
|
||||
this.searchRadius(far, query, radius, results, depth + 1);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Private: collect ───────────────────────────────────────────────────────
|
||||
|
||||
private collect(node: KDNode<T> | null, out: KDPoint<T>[]): void {
|
||||
if (node === null) return;
|
||||
out.push(node.point);
|
||||
this.collect(node.left, out);
|
||||
this.collect(node.right, out);
|
||||
}
|
||||
|
||||
// ── Private: validation ────────────────────────────────────────────────────
|
||||
|
||||
private validateVector(v: number[]): void {
|
||||
if (v.length !== this.dims) {
|
||||
throw new TypeError(
|
||||
`Vector length ${v.length} does not match tree dimensionality ${this.dims}`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private validate(point: KDPoint<T>): void {
|
||||
this.validateVector(point.vector);
|
||||
}
|
||||
|
||||
private validateAll(points: KDPoint<T>[]): void {
|
||||
for (const p of points) this.validate(p);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user