mirror of
https://github.com/abhigyanpatwari/GitNexus.git
synced 2026-10-13 03:58:33 +00:00
fix(embeddings): spill cached vectors to a Float32 temp file
Keep restore metadata in RAM and write embeddings once the in-memory row limit is exceeded so incremental analyze can survive large caches without a full-table number[] heap (#3306). Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Gergo Magyar
parent
d2a43e33df
commit
a5ad6faeee
2 changed files with 558 additions and 0 deletions
385
gitnexus/src/core/embeddings/embedding-restore-spill.ts
Normal file
385
gitnexus/src/core/embeddings/embedding-restore-spill.ts
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
/**
|
||||
* Disk-backed restore cache for CodeEmbedding rows (#3306).
|
||||
*
|
||||
* `loadCachedEmbeddings` used to `getAll()` the table and `map(Number)` every
|
||||
* vector into a JS `number[]`. On a large already-indexed repo that single
|
||||
* structure OOMs the V8 heap during "Caching embeddings..." even when the
|
||||
* incremental diff is a handful of nodes.
|
||||
*
|
||||
* This module keeps metadata in RAM and writes vectors to a temp Float32
|
||||
* spill. Restore materializes only the rows that Phase 3.5 will re-insert,
|
||||
* in the existing 200-row batches.
|
||||
*/
|
||||
import { closeSync, existsSync, openSync, readSync, unlinkSync, writeSync } from 'node:fs';
|
||||
import { randomBytes } from 'node:crypto';
|
||||
import os from 'node:os';
|
||||
import path from 'node:path';
|
||||
import type { CachedEmbedding } from './types.js';
|
||||
|
||||
/** In-RAM vector copies stay below this row count; larger tables use the spill. */
|
||||
export const DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT = 2048;
|
||||
|
||||
const SPILL_MAGIC = 'GNXE';
|
||||
const SPILL_VERSION = 1;
|
||||
const SPILL_HEADER_BYTES = 12;
|
||||
|
||||
export interface CachedEmbeddingMeta {
|
||||
nodeId: string;
|
||||
chunkIndex: number;
|
||||
startLine: number;
|
||||
endLine: number;
|
||||
contentHash?: string;
|
||||
/** Row order in the spill file (and in `embeddings` when in-memory). */
|
||||
vectorIndex: number;
|
||||
}
|
||||
|
||||
export interface EmbeddingVectorSpill {
|
||||
path: string;
|
||||
dims: number;
|
||||
rowCount: number;
|
||||
}
|
||||
|
||||
export interface CachedEmbeddingsSnapshot {
|
||||
embeddingNodeIds: Set<string>;
|
||||
/** Populated only when the table is at or under the in-memory row limit. */
|
||||
embeddings: CachedEmbedding[];
|
||||
rows: CachedEmbeddingMeta[];
|
||||
spill?: EmbeddingVectorSpill;
|
||||
}
|
||||
|
||||
export interface LoadCachedEmbeddingsOptions {
|
||||
/**
|
||||
* Keep full `number[]` vectors in RAM at or below this many rows.
|
||||
* `0` always spills. Default {@link DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT}
|
||||
* or `GITNEXUS_EMBEDDING_CACHE_IN_MEMORY_LIMIT`.
|
||||
*/
|
||||
inMemoryRowLimit?: number;
|
||||
/** Directory for the spill file (default `os.tmpdir()`). */
|
||||
spillDir?: string;
|
||||
}
|
||||
|
||||
export interface CachedEmbeddingsBuilder {
|
||||
embeddingNodeIds: Set<string>;
|
||||
rows: CachedEmbeddingMeta[];
|
||||
inMemory: CachedEmbedding[] | null;
|
||||
inMemoryRowLimit: number;
|
||||
writer: EmbeddingSpillWriter;
|
||||
}
|
||||
|
||||
export function emptyCachedEmbeddingsSnapshot(): CachedEmbeddingsSnapshot {
|
||||
return { embeddingNodeIds: new Set(), embeddings: [], rows: [] };
|
||||
}
|
||||
|
||||
export function resolveEmbeddingCacheInMemoryRowLimit(override?: number): number {
|
||||
if (override !== undefined) {
|
||||
if (!Number.isFinite(override) || override < 0) {
|
||||
return DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT;
|
||||
}
|
||||
return Math.floor(override);
|
||||
}
|
||||
const raw = process.env.GITNEXUS_EMBEDDING_CACHE_IN_MEMORY_LIMIT;
|
||||
if (raw === undefined || raw === '') return DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT;
|
||||
const parsed = parseInt(raw, 10);
|
||||
return Number.isFinite(parsed) && parsed >= 0
|
||||
? parsed
|
||||
: DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT;
|
||||
}
|
||||
|
||||
export function normalizeCachedEmbeddings(raw: {
|
||||
embeddingNodeIds?: Set<string>;
|
||||
embeddings?: CachedEmbedding[];
|
||||
rows?: CachedEmbeddingMeta[];
|
||||
spill?: EmbeddingVectorSpill;
|
||||
}): CachedEmbeddingsSnapshot {
|
||||
const embeddings = raw.embeddings ?? [];
|
||||
const embeddingNodeIds = raw.embeddingNodeIds ?? new Set(embeddings.map((row) => row.nodeId));
|
||||
const rows =
|
||||
raw.rows ??
|
||||
embeddings.map((row, vectorIndex) => ({
|
||||
nodeId: row.nodeId,
|
||||
chunkIndex: row.chunkIndex,
|
||||
startLine: row.startLine,
|
||||
endLine: row.endLine,
|
||||
contentHash: row.contentHash,
|
||||
vectorIndex,
|
||||
}));
|
||||
return { embeddingNodeIds, embeddings, rows, spill: raw.spill };
|
||||
}
|
||||
|
||||
export function cacheRowCount(snapshot: CachedEmbeddingsSnapshot): number {
|
||||
return snapshot.rows.length > 0 ? snapshot.rows.length : snapshot.embeddings.length;
|
||||
}
|
||||
|
||||
export function snapshotEmbeddingDims(snapshot: CachedEmbeddingsSnapshot): number | undefined {
|
||||
if (snapshot.spill && snapshot.spill.dims > 0) return snapshot.spill.dims;
|
||||
const dims = snapshot.embeddings[0]?.embedding.length;
|
||||
return dims && dims > 0 ? dims : undefined;
|
||||
}
|
||||
|
||||
export function coerceEmbeddingToFloat32(embedding: unknown): Float32Array | null {
|
||||
if (embedding == null) return null;
|
||||
if (embedding instanceof Float32Array) {
|
||||
return embedding.length > 0 ? embedding : null;
|
||||
}
|
||||
if (ArrayBuffer.isView(embedding) && !(embedding instanceof DataView)) {
|
||||
const view = embedding as Exclude<ArrayBufferView, DataView> & { length: number };
|
||||
if (view.length === 0) return null;
|
||||
const out = new Float32Array(view.length);
|
||||
for (let i = 0; i < view.length; i++) out[i] = Number(view[i]);
|
||||
return out;
|
||||
}
|
||||
if (
|
||||
typeof embedding === 'object' &&
|
||||
typeof (embedding as Iterable<unknown>)[Symbol.iterator] === 'function'
|
||||
) {
|
||||
const arr = Array.isArray(embedding)
|
||||
? (embedding as unknown[])
|
||||
: Array.from(embedding as Iterable<unknown>);
|
||||
if (arr.length === 0) return null;
|
||||
const out = new Float32Array(arr.length);
|
||||
for (let i = 0; i < arr.length; i++) out[i] = Number(arr[i]);
|
||||
return out;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
export function float32ToNumberArray(vec: Float32Array): number[] {
|
||||
const out = new Array<number>(vec.length);
|
||||
for (let i = 0; i < vec.length; i++) out[i] = vec[i]!;
|
||||
return out;
|
||||
}
|
||||
|
||||
export class EmbeddingSpillWriter {
|
||||
readonly path: string;
|
||||
dims = 0;
|
||||
rowCount = 0;
|
||||
private fd: number | null = null;
|
||||
private closed = false;
|
||||
|
||||
constructor(dir: string) {
|
||||
this.path = path.join(
|
||||
dir,
|
||||
`gitnexus-embed-restore-${process.pid}-${randomBytes(8).toString('hex')}.bin`,
|
||||
);
|
||||
}
|
||||
|
||||
append(vec: Float32Array): void {
|
||||
if (this.closed) {
|
||||
throw new Error('embedding spill writer already closed');
|
||||
}
|
||||
if (this.fd === null) {
|
||||
this.dims = vec.length;
|
||||
this.fd = openSync(this.path, 'wx', 0o600);
|
||||
const header = Buffer.alloc(SPILL_HEADER_BYTES);
|
||||
header.write(SPILL_MAGIC, 0, 4, 'ascii');
|
||||
header.writeUInt8(SPILL_VERSION, 4);
|
||||
header.writeUInt32LE(this.dims, 5);
|
||||
writeSync(this.fd, header);
|
||||
} else if (vec.length !== this.dims) {
|
||||
throw new Error(
|
||||
`embedding dim mismatch while spilling: got ${vec.length}, expected ${this.dims}`,
|
||||
);
|
||||
}
|
||||
writeSync(this.fd, Buffer.from(vec.buffer, vec.byteOffset, vec.byteLength));
|
||||
this.rowCount++;
|
||||
}
|
||||
|
||||
finish(): EmbeddingVectorSpill | undefined {
|
||||
if (this.fd !== null) {
|
||||
closeSync(this.fd);
|
||||
this.fd = null;
|
||||
}
|
||||
this.closed = true;
|
||||
if (this.rowCount === 0) {
|
||||
this.unlinkQuiet();
|
||||
return undefined;
|
||||
}
|
||||
return { path: this.path, dims: this.dims, rowCount: this.rowCount };
|
||||
}
|
||||
|
||||
abort(): void {
|
||||
if (this.fd !== null) {
|
||||
try {
|
||||
closeSync(this.fd);
|
||||
} catch {
|
||||
/* already closed */
|
||||
}
|
||||
this.fd = null;
|
||||
}
|
||||
this.closed = true;
|
||||
this.unlinkQuiet();
|
||||
}
|
||||
|
||||
private unlinkQuiet(): void {
|
||||
try {
|
||||
unlinkSync(this.path);
|
||||
} catch {
|
||||
/* ENOENT or already removed */
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function readSpillVectors(
|
||||
spill: EmbeddingVectorSpill,
|
||||
indices: readonly number[],
|
||||
): Float32Array[] {
|
||||
const fd = openSync(spill.path, 'r');
|
||||
try {
|
||||
const header = Buffer.alloc(SPILL_HEADER_BYTES);
|
||||
const headerRead = readSync(fd, header, 0, SPILL_HEADER_BYTES, 0);
|
||||
if (headerRead !== SPILL_HEADER_BYTES || header.toString('ascii', 0, 4) !== SPILL_MAGIC) {
|
||||
throw new Error(`invalid embedding spill header: ${spill.path}`);
|
||||
}
|
||||
if (header.readUInt8(4) !== SPILL_VERSION) {
|
||||
throw new Error(`unsupported embedding spill version in ${spill.path}`);
|
||||
}
|
||||
const dims = header.readUInt32LE(5);
|
||||
if (dims !== spill.dims) {
|
||||
throw new Error(`embedding spill dim mismatch: file ${dims}, expected ${spill.dims}`);
|
||||
}
|
||||
const bytesPerVec = dims * 4;
|
||||
const out: Float32Array[] = [];
|
||||
const buf = Buffer.alloc(bytesPerVec);
|
||||
for (const index of indices) {
|
||||
if (!Number.isInteger(index) || index < 0 || index >= spill.rowCount) {
|
||||
throw new Error(`embedding spill index out of range: ${index}`);
|
||||
}
|
||||
const offset = SPILL_HEADER_BYTES + index * bytesPerVec;
|
||||
const n = readSync(fd, buf, 0, bytesPerVec, offset);
|
||||
if (n !== bytesPerVec) {
|
||||
throw new Error(`short embedding spill read at index ${index}`);
|
||||
}
|
||||
const copy = new Float32Array(dims);
|
||||
Buffer.from(copy.buffer).set(buf);
|
||||
out.push(copy);
|
||||
}
|
||||
return out;
|
||||
} finally {
|
||||
closeSync(fd);
|
||||
}
|
||||
}
|
||||
|
||||
export function disposeEmbeddingSpill(spill?: EmbeddingVectorSpill): void {
|
||||
if (!spill?.path) return;
|
||||
try {
|
||||
if (existsSync(spill.path)) unlinkSync(spill.path);
|
||||
} catch {
|
||||
/* best-effort */
|
||||
}
|
||||
}
|
||||
|
||||
export function createCachedEmbeddingsBuilder(
|
||||
options?: LoadCachedEmbeddingsOptions,
|
||||
): CachedEmbeddingsBuilder {
|
||||
const inMemoryRowLimit = resolveEmbeddingCacheInMemoryRowLimit(options?.inMemoryRowLimit);
|
||||
return {
|
||||
embeddingNodeIds: new Set(),
|
||||
rows: [],
|
||||
inMemory: inMemoryRowLimit <= 0 ? null : [],
|
||||
inMemoryRowLimit,
|
||||
writer: new EmbeddingSpillWriter(options?.spillDir ?? os.tmpdir()),
|
||||
};
|
||||
}
|
||||
|
||||
export function ingestCachedEmbeddingRow(
|
||||
builder: CachedEmbeddingsBuilder,
|
||||
row: Record<string, unknown> | unknown[],
|
||||
hasContentHash: boolean,
|
||||
): void {
|
||||
const rec = row as Record<string, unknown> & unknown[];
|
||||
const nodeId = String(rec.nodeId ?? rec[0] ?? '');
|
||||
if (!nodeId) return;
|
||||
const embedding = rec.embedding ?? rec[4];
|
||||
const f32 = coerceEmbeddingToFloat32(embedding);
|
||||
if (!f32) return;
|
||||
|
||||
builder.embeddingNodeIds.add(nodeId);
|
||||
const meta: CachedEmbeddingMeta = {
|
||||
nodeId,
|
||||
chunkIndex: Number(rec.chunkIndex ?? rec[1] ?? 0),
|
||||
startLine: Number(rec.startLine ?? rec[2] ?? 0),
|
||||
endLine: Number(rec.endLine ?? rec[3] ?? 0),
|
||||
contentHash: hasContentHash
|
||||
? ((rec.contentHash ?? rec[5] ?? undefined) as string | undefined)
|
||||
: undefined,
|
||||
vectorIndex: builder.rows.length,
|
||||
};
|
||||
builder.rows.push(meta);
|
||||
|
||||
if (builder.inMemory && builder.rows.length <= builder.inMemoryRowLimit) {
|
||||
builder.inMemory.push({
|
||||
nodeId: meta.nodeId,
|
||||
chunkIndex: meta.chunkIndex,
|
||||
startLine: meta.startLine,
|
||||
endLine: meta.endLine,
|
||||
contentHash: meta.contentHash,
|
||||
embedding: float32ToNumberArray(f32),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (builder.inMemory) {
|
||||
for (const cached of builder.inMemory) {
|
||||
const prior = coerceEmbeddingToFloat32(cached.embedding);
|
||||
if (prior) builder.writer.append(prior);
|
||||
}
|
||||
builder.inMemory = null;
|
||||
}
|
||||
builder.writer.append(f32);
|
||||
}
|
||||
|
||||
export function finalizeCachedEmbeddingsSnapshot(
|
||||
builder: CachedEmbeddingsBuilder,
|
||||
): CachedEmbeddingsSnapshot {
|
||||
if (builder.inMemory) {
|
||||
builder.writer.abort();
|
||||
return {
|
||||
embeddingNodeIds: builder.embeddingNodeIds,
|
||||
embeddings: builder.inMemory,
|
||||
rows: builder.rows,
|
||||
};
|
||||
}
|
||||
return {
|
||||
embeddingNodeIds: builder.embeddingNodeIds,
|
||||
embeddings: [],
|
||||
rows: builder.rows,
|
||||
spill: builder.writer.finish(),
|
||||
};
|
||||
}
|
||||
|
||||
export function abortCachedEmbeddingsBuilder(builder: CachedEmbeddingsBuilder): void {
|
||||
builder.writer.abort();
|
||||
}
|
||||
|
||||
export function materializeCachedEmbeddings(
|
||||
snapshot: CachedEmbeddingsSnapshot,
|
||||
metas: readonly CachedEmbeddingMeta[],
|
||||
): CachedEmbedding[] {
|
||||
if (metas.length === 0) return [];
|
||||
if (snapshot.spill && snapshot.embeddings.length === 0) {
|
||||
const vectors = readSpillVectors(
|
||||
snapshot.spill,
|
||||
metas.map((meta) => meta.vectorIndex),
|
||||
);
|
||||
return metas.map((meta, i) => ({
|
||||
nodeId: meta.nodeId,
|
||||
chunkIndex: meta.chunkIndex,
|
||||
startLine: meta.startLine,
|
||||
endLine: meta.endLine,
|
||||
contentHash: meta.contentHash,
|
||||
embedding: float32ToNumberArray(vectors[i]!),
|
||||
}));
|
||||
}
|
||||
if (snapshot.embeddings.length === 0) return [];
|
||||
const byKey = new Map(
|
||||
snapshot.embeddings.map((row) => [`${row.nodeId}:${row.chunkIndex}`, row] as const),
|
||||
);
|
||||
return metas.map((meta) => {
|
||||
const hit =
|
||||
byKey.get(`${meta.nodeId}:${meta.chunkIndex}`) ?? snapshot.embeddings[meta.vectorIndex];
|
||||
if (!hit) {
|
||||
throw new Error(`missing cached embedding ${meta.nodeId}:${meta.chunkIndex}`);
|
||||
}
|
||||
return hit;
|
||||
});
|
||||
}
|
||||
173
gitnexus/test/unit/embedding-restore-spill.test.ts
Normal file
173
gitnexus/test/unit/embedding-restore-spill.test.ts
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
import { existsSync, statSync, writeFileSync } from 'node:fs';
|
||||
import os from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest';
|
||||
import {
|
||||
abortCachedEmbeddingsBuilder,
|
||||
cacheRowCount,
|
||||
createCachedEmbeddingsBuilder,
|
||||
DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT,
|
||||
disposeEmbeddingSpill,
|
||||
finalizeCachedEmbeddingsSnapshot,
|
||||
ingestCachedEmbeddingRow,
|
||||
materializeCachedEmbeddings,
|
||||
normalizeCachedEmbeddings,
|
||||
readSpillVectors,
|
||||
resolveEmbeddingCacheInMemoryRowLimit,
|
||||
snapshotEmbeddingDims,
|
||||
} from '../../src/core/embeddings/embedding-restore-spill.js';
|
||||
|
||||
const DIMS = 8;
|
||||
|
||||
function vector(fill: number): number[] {
|
||||
return Array.from({ length: DIMS }, () => fill);
|
||||
}
|
||||
|
||||
function row(id: string, fill: number, hash = `hash-${id}`) {
|
||||
return {
|
||||
nodeId: id,
|
||||
chunkIndex: 0,
|
||||
startLine: 1,
|
||||
endLine: 2,
|
||||
embedding: vector(fill),
|
||||
contentHash: hash,
|
||||
};
|
||||
}
|
||||
|
||||
describe('embedding-restore-spill (#3306)', () => {
|
||||
const spills: Array<{ path: string }> = [];
|
||||
afterEach(() => {
|
||||
for (const spill of spills) disposeEmbeddingSpill(spill);
|
||||
spills.length = 0;
|
||||
});
|
||||
|
||||
it('keeps small tables in RAM and does not leave a spill file', () => {
|
||||
const builder = createCachedEmbeddingsBuilder({
|
||||
inMemoryRowLimit: 4,
|
||||
spillDir: os.tmpdir(),
|
||||
});
|
||||
ingestCachedEmbeddingRow(builder, row('n1', 0.25), true);
|
||||
ingestCachedEmbeddingRow(builder, row('n2', 0.5), true);
|
||||
const snapshot = finalizeCachedEmbeddingsSnapshot(builder);
|
||||
expect(snapshot.spill).toBeUndefined();
|
||||
expect(existsSync(builder.writer.path)).toBe(false);
|
||||
expect(snapshot.embeddings).toHaveLength(2);
|
||||
expect(snapshot.rows).toHaveLength(2);
|
||||
expect(snapshot.embeddings[0]?.embedding[0]).toBeCloseTo(0.25);
|
||||
expect(cacheRowCount(snapshot)).toBe(2);
|
||||
expect(snapshotEmbeddingDims(snapshot)).toBe(DIMS);
|
||||
});
|
||||
|
||||
it('spills vectors once the in-memory limit is exceeded and materializes a subset', () => {
|
||||
const builder = createCachedEmbeddingsBuilder({
|
||||
inMemoryRowLimit: 2,
|
||||
spillDir: os.tmpdir(),
|
||||
});
|
||||
for (let i = 0; i < 5; i++) {
|
||||
ingestCachedEmbeddingRow(builder, row(`n${i}`, i + 1), true);
|
||||
}
|
||||
const snapshot = finalizeCachedEmbeddingsSnapshot(builder);
|
||||
if (snapshot.spill) spills.push(snapshot.spill);
|
||||
expect(snapshot.embeddings).toEqual([]);
|
||||
expect(snapshot.spill?.rowCount).toBe(5);
|
||||
expect(existsSync(snapshot.spill!.path)).toBe(true);
|
||||
expect(snapshot.embeddingNodeIds.size).toBe(5);
|
||||
|
||||
const subset = materializeCachedEmbeddings(snapshot, snapshot.rows.slice(1, 3));
|
||||
expect(subset).toHaveLength(2);
|
||||
expect(subset[0]?.nodeId).toBe('n1');
|
||||
expect(subset[0]?.embedding[0]).toBeCloseTo(2);
|
||||
expect(subset[1]?.embedding[0]).toBeCloseTo(3);
|
||||
});
|
||||
|
||||
it('always spills when the in-memory limit is 0 (no Number[] table in RAM)', () => {
|
||||
const builder = createCachedEmbeddingsBuilder({
|
||||
inMemoryRowLimit: 0,
|
||||
spillDir: os.tmpdir(),
|
||||
});
|
||||
ingestCachedEmbeddingRow(builder, row('only', 0.75), true);
|
||||
const snapshot = finalizeCachedEmbeddingsSnapshot(builder);
|
||||
if (snapshot.spill) spills.push(snapshot.spill);
|
||||
expect(snapshot.embeddings).toEqual([]);
|
||||
expect(snapshot.rows).toHaveLength(1);
|
||||
expect(materializeCachedEmbeddings(snapshot, snapshot.rows)[0]?.embedding[0]).toBeCloseTo(0.75);
|
||||
});
|
||||
|
||||
it('aborts an unfinished builder without leaking a spill file', () => {
|
||||
const builder = createCachedEmbeddingsBuilder({
|
||||
inMemoryRowLimit: 0,
|
||||
spillDir: os.tmpdir(),
|
||||
});
|
||||
ingestCachedEmbeddingRow(builder, row('n1', 1), true);
|
||||
expect(existsSync(builder.writer.path)).toBe(true);
|
||||
abortCachedEmbeddingsBuilder(builder);
|
||||
expect(existsSync(builder.writer.path)).toBe(false);
|
||||
});
|
||||
|
||||
it('normalizes mock {embeddings} payloads so Phase 3.5 can restore without a spill', () => {
|
||||
const snapshot = normalizeCachedEmbeddings({
|
||||
embeddingNodeIds: new Set(['Function:a:foo']),
|
||||
embeddings: [
|
||||
{
|
||||
nodeId: 'Function:a:foo',
|
||||
chunkIndex: 0,
|
||||
startLine: 0,
|
||||
endLine: 3,
|
||||
embedding: vector(0.1),
|
||||
contentHash: 'stub',
|
||||
},
|
||||
],
|
||||
});
|
||||
expect(snapshot.rows).toHaveLength(1);
|
||||
expect(snapshot.rows[0]?.vectorIndex).toBe(0);
|
||||
const restored = materializeCachedEmbeddings(snapshot, snapshot.rows);
|
||||
expect(restored[0]?.contentHash).toBe('stub');
|
||||
expect(restored[0]?.embedding).toHaveLength(DIMS);
|
||||
});
|
||||
|
||||
it('defaults the in-memory row limit to 2048 and honors GITNEXUS_EMBEDDING_CACHE_IN_MEMORY_LIMIT', () => {
|
||||
expect(DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT).toBe(2048);
|
||||
expect(resolveEmbeddingCacheInMemoryRowLimit()).toBe(2048);
|
||||
vi.stubEnv('GITNEXUS_EMBEDDING_CACHE_IN_MEMORY_LIMIT', '0');
|
||||
expect(resolveEmbeddingCacheInMemoryRowLimit()).toBe(0);
|
||||
vi.stubEnv('GITNEXUS_EMBEDDING_CACHE_IN_MEMORY_LIMIT', '12');
|
||||
expect(resolveEmbeddingCacheInMemoryRowLimit()).toBe(12);
|
||||
vi.stubEnv('GITNEXUS_EMBEDDING_CACHE_IN_MEMORY_LIMIT', 'nope');
|
||||
expect(resolveEmbeddingCacheInMemoryRowLimit()).toBe(2048);
|
||||
vi.unstubAllEnvs();
|
||||
expect(resolveEmbeddingCacheInMemoryRowLimit(7)).toBe(7);
|
||||
expect(resolveEmbeddingCacheInMemoryRowLimit(-1)).toBe(2048);
|
||||
});
|
||||
|
||||
it('spills above the default 2048-row limit with header-plus-body size 12 + N * D * 4', () => {
|
||||
const n = DEFAULT_EMBEDDING_CACHE_IN_MEMORY_ROW_LIMIT + 1;
|
||||
const builder = createCachedEmbeddingsBuilder({ spillDir: os.tmpdir() });
|
||||
for (let i = 0; i < n; i++) {
|
||||
ingestCachedEmbeddingRow(builder, row(`n${i}`, 1), true);
|
||||
}
|
||||
const snapshot = finalizeCachedEmbeddingsSnapshot(builder);
|
||||
if (snapshot.spill) spills.push(snapshot.spill);
|
||||
expect(snapshot.embeddings).toEqual([]);
|
||||
expect(snapshot.spill?.rowCount).toBe(n);
|
||||
expect(statSync(snapshot.spill!.path).size).toBe(12 + n * DIMS * 4);
|
||||
});
|
||||
|
||||
it('rejects a dim mismatch once spilling and rejects a short or bad-magic header', () => {
|
||||
const builder = createCachedEmbeddingsBuilder({
|
||||
inMemoryRowLimit: 0,
|
||||
spillDir: os.tmpdir(),
|
||||
});
|
||||
ingestCachedEmbeddingRow(builder, row('a', 1), true);
|
||||
expect(() =>
|
||||
ingestCachedEmbeddingRow(builder, { ...row('b', 2), embedding: [1, 2, 3] }, true),
|
||||
).toThrow(/dim mismatch/);
|
||||
abortCachedEmbeddingsBuilder(builder);
|
||||
|
||||
const badPath = path.join(os.tmpdir(), `gitnexus-embed-restore-bad-${process.pid}.bin`);
|
||||
writeFileSync(badPath, Buffer.from('NOPE'));
|
||||
spills.push({ path: badPath });
|
||||
expect(() => readSpillVectors({ path: badPath, dims: DIMS, rowCount: 1 }, [0])).toThrow(
|
||||
/invalid embedding spill header/,
|
||||
);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Reference in a new issue