diff --git a/README.md b/README.md index b48c92b829..05bfda25c1 100644 --- a/README.md +++ b/README.md @@ -575,6 +575,97 @@ Give the LLM the schema, a few examples, and let it compose queries. The schema --- +## ๐Ÿ”ฌ Deep Dive: Copy-on-Write Woes with In-Memory WASM Databases + +While building the embedding pipeline, we hit an interesting memory problem. Documenting it here because it's a non-obvious gotcha for anyone doing vector storage in browser-side databases. + +### The Setup + +We wanted to store 384-dimensional embeddings alongside our code nodes. Natural instinct: add an `embedding FLOAT[384]` column to the existing `CodeNode` table, bulk load the graph, then `UPDATE` each node with its embedding. + +```cypher +-- Seemed reasonable, right? +MATCH (n:CodeNode {id: $id}) SET n.embedding = $vec +``` + +### The Problem + +Worked fine for ~20 nodes. Exploded at ~1000 nodes with: + +``` +Buffer manager exception: Unable to allocate memory! The buffer pool is full! +``` + +We had a 512MB buffer pool. 1000 embeddings ร— 384 floats ร— 4 bytes = ~1.5MB. Where did 512MB go? + +**Answer: Copy-on-Write (COW).** + +Most databases don't modify records in place. When you `UPDATE`, they create a new version of the record (for transaction rollback, MVCC, etc.). The old version sticks around until commit. + +Our `CodeNode` table had a `content` field averaging ~2KB per node (code snippets). So each `UPDATE`: + +1. Reads the entire node (~2KB) +2. Creates a new copy with the embedding (~3.5KB) +3. Keeps the old version around + +For 1000 nodes: `1000 ร— 2KB (old) + 1000 ร— 3.5KB (new) = ~5.5MB`... but that's just user data. KuzuDB's internal structures (indexes, hash tables, page management) multiply this significantly. And since it's an in-memory database, the buffer pool IS the storage - there's no disk to spill to. + +```mermaid +flowchart LR + subgraph Before["Before UPDATE"] + N1[CodeNode
id + name + content
~2KB] + end + + subgraph During["During UPDATE (COW)"] + N1_OLD[Old Version
~2KB] + N1_NEW[New Version
+ embedding
~3.5KB] + end + + subgraph Problem["ร— 1000 nodes"] + BOOM[๐Ÿ’ฅ Buffer Pool Exhausted] + end + + Before --> During --> Problem +``` + +### The Fix: Separate Table Architecture + +Don't `UPDATE` wide tables. `INSERT` into a narrow one. + +```mermaid +flowchart TD + subgraph Old["โŒ Original Design"] + CN1[CodeNode
id, name, content, embedding
~3.5KB per UPDATE copy] + end + + subgraph New["โœ… New Design"] + CN2[CodeNode
id, name, content] + CE[CodeEmbedding
nodeId, embedding
~1.5KB INSERT only] + end + + Old -->|"COW copies entire 2KB+ node"| FAIL[Memory Explosion] + New -->|"INSERT into lightweight table"| WIN[Works at scale] +``` + +Now we: +1. Bulk load `CodeNode` (no embedding column) +2. `CREATE` rows in `CodeEmbedding` table (just `nodeId` + `embedding`) +3. Vector index lives on `CodeEmbedding` +4. Semantic search JOINs back to `CodeNode` for metadata + +**Trade-off:** Every semantic search needs a JOIN. But it's a primary key lookup (O(1)), so we're talking ~1-5ms extra per query. Totally worth it to not explode at 1000 nodes. + +### Lessons Learned + +1. **In-memory WASM DBs have hard limits** - No disk spillover, buffer pool is everything +2. **COW amplifies record size** - That innocent `UPDATE` copies your whole row +3. **Normalize for bulk writes** - Especially for append-only data like embeddings +4. **Profile the pathological case** - 20 nodes worked, 1000 didn't. Always test at scale + +This is one of those "obvious in hindsight" things. Most vector DB tutorials show single-table schemas because they're using databases with disk backing. In-browser WASM land plays by different rules. + +--- + ## Security & Privacy - All processing happens in your browser diff --git a/src/components/EmbeddingStatus.tsx b/src/components/EmbeddingStatus.tsx index 700d0dd019..f7e2bed2a8 100644 --- a/src/components/EmbeddingStatus.tsx +++ b/src/components/EmbeddingStatus.tsx @@ -1,5 +1,6 @@ -import { Brain, Loader2, Check, AlertCircle, Zap } from 'lucide-react'; +import { Brain, Loader2, Check, AlertCircle, Zap, FlaskConical } from 'lucide-react'; import { useAppState } from '../hooks/useAppState'; +import { useState } from 'react'; /** * Embedding status indicator and trigger button @@ -11,8 +12,11 @@ export const EmbeddingStatus = () => { embeddingProgress, startEmbeddings, graph, - viewMode + viewMode, + testArrayParams, } = useAppState(); + + const [testResult, setTestResult] = useState(null); // Only show when exploring a loaded graph if (viewMode !== 'exploring' || !graph) return null; @@ -24,19 +28,45 @@ export const EmbeddingStatus = () => { console.error('Embedding failed:', error); } }; + + const handleTestArrayParams = async () => { + setTestResult('Testing...'); + const result = await testArrayParams(); + if (result.success) { + setTestResult('โœ… Array params WORK!'); + console.log('โœ… Array params test passed!'); + } else { + setTestResult(`โŒ ${result.error}`); + console.error('โŒ Array params test failed:', result.error); + } + }; // Idle state - show button to start if (embeddingStatus === 'idle') { return ( - +
+ {/* Test button (dev only) */} + {import.meta.env.DEV && ( + + )} + + +
); } diff --git a/src/core/embeddings/embedding-pipeline.ts b/src/core/embeddings/embedding-pipeline.ts index 6b38c2b241..6a0d8619a6 100644 --- a/src/core/embeddings/embedding-pipeline.ts +++ b/src/core/embeddings/embedding-pipeline.ts @@ -59,47 +59,32 @@ const queryEmbeddableNodes = async ( }; /** - * Update a single node's embedding in KuzuDB + * Batch INSERT embeddings into separate CodeEmbedding table + * Using a separate lightweight table avoids copy-on-write overhead + * that occurs when UPDATEing nodes with large content fields */ -const updateNodeEmbedding = async ( - executeQuery: (cypher: string) => Promise, - nodeId: string, - embedding: number[] -): Promise => { - // KuzuDB requires the array to be cast to the correct type - const embeddingStr = `[${embedding.join(',')}]`; - - const cypher = ` - MATCH (n:CodeNode {id: '${nodeId}'}) - SET n.embedding = CAST(${embeddingStr} AS FLOAT[384]) - `; - - await executeQuery(cypher); -}; - -/** - * Batch update multiple node embeddings - * More efficient than individual updates - */ -const batchUpdateEmbeddings = async ( - executeQuery: (cypher: string) => Promise, +const batchInsertEmbeddings = async ( + executeWithReusedStatement: ( + cypher: string, + paramsList: Array> + ) => Promise, updates: Array<{ id: string; embedding: number[] }> ): Promise => { - // Process updates one by one for now - // KuzuDB doesn't have great batch update syntax - for (const update of updates) { - await updateNodeEmbedding(executeQuery, update.id, update.embedding); - } + // INSERT into separate embedding table - much more memory efficient! + const cypher = `CREATE (e:CodeEmbedding {nodeId: $nodeId, embedding: $embedding})`; + const paramsList = updates.map(u => ({ nodeId: u.id, embedding: u.embedding })); + await executeWithReusedStatement(cypher, paramsList); }; /** * Create the vector index for semantic search + * Now indexes the separate CodeEmbedding table */ const createVectorIndex = async ( executeQuery: (cypher: string) => Promise ): Promise => { const cypher = ` - CALL CREATE_VECTOR_INDEX('CodeNode', 'code_embedding_idx', 'embedding', metric := 'cosine') + CALL CREATE_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx', 'embedding', metric := 'cosine') `; try { @@ -116,11 +101,13 @@ const createVectorIndex = async ( * Run the embedding pipeline * * @param executeQuery - Function to execute Cypher queries against KuzuDB + * @param executeWithReusedStatement - Function to execute with reused prepared statement * @param onProgress - Callback for progress updates * @param config - Optional configuration override */ export const runEmbeddingPipeline = async ( executeQuery: (cypher: string) => Promise, + executeWithReusedStatement: (cypher: string, paramsList: Array>) => Promise, onProgress: EmbeddingProgressCallback, config: Partial = {} ): Promise => { @@ -203,7 +190,7 @@ export const runEmbeddingPipeline = async ( embedding: embeddingToArray(embeddings[i]), })); - await batchUpdateEmbeddings(executeQuery, updates); + await batchInsertEmbeddings(executeWithReusedStatement, updates); processedNodes += batch.length; @@ -264,6 +251,8 @@ export const runEmbeddingPipeline = async ( /** * Perform semantic search using the vector index * + * Uses separate CodeEmbedding table and JOINs with CodeNode for metadata + * * @param executeQuery - Function to execute Cypher queries * @param query - Search query text * @param k - Number of results to return (default: 10) @@ -285,15 +274,16 @@ export const semanticSearch = async ( const queryVec = embeddingToArray(queryEmbedding); const queryVecStr = `[${queryVec.join(',')}]`; - // Query the vector index + // Query the vector index on CodeEmbedding, then JOIN with CodeNode for metadata const cypher = ` - CALL QUERY_VECTOR_INDEX('CodeNode', 'code_embedding_idx', + CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx', CAST(${queryVecStr} AS FLOAT[384]), ${k}) - YIELD node, distance + YIELD node AS emb, distance WHERE distance < ${maxDistance} - RETURN node.id AS nodeId, node.name AS name, node.label AS label, - node.filePath AS filePath, distance, - node.startLine AS startLine, node.endLine AS endLine + MATCH (n:CodeNode {id: emb.nodeId}) + RETURN n.id AS nodeId, n.name AS name, n.label AS label, + n.filePath AS filePath, distance, + n.startLine AS startLine, n.endLine AS endLine ORDER BY distance `; @@ -314,6 +304,8 @@ export const semanticSearch = async ( * Semantic search with graph expansion * Finds similar nodes AND their connections * + * Uses separate CodeEmbedding table and JOINs with CodeNode + * * @param executeQuery - Function to execute Cypher queries * @param query - Search query text * @param k - Number of initial results @@ -335,12 +327,13 @@ export const semanticSearchWithContext = async ( const queryVec = embeddingToArray(queryEmbedding); const queryVecStr = `[${queryVec.join(',')}]`; - // Query with graph expansion + // Query embedding table, JOIN with CodeNode, then expand graph const cypher = ` - CALL QUERY_VECTOR_INDEX('CodeNode', 'code_embedding_idx', + CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx', CAST(${queryVecStr} AS FLOAT[384]), ${k}) - YIELD node AS match, distance + YIELD node AS emb, distance WHERE distance < 0.5 + MATCH (match:CodeNode {id: emb.nodeId}) MATCH (match)-[r:CodeRelation*1..${hops}]-(connected:CodeNode) RETURN match.id AS matchId, match.name AS matchName, match.label AS matchLabel, match.filePath AS matchPath, distance, diff --git a/src/core/kuzu/kuzu-adapter.ts b/src/core/kuzu/kuzu-adapter.ts index 59710e0344..eea3710d54 100644 --- a/src/core/kuzu/kuzu-adapter.ts +++ b/src/core/kuzu/kuzu-adapter.ts @@ -8,7 +8,7 @@ */ import { KnowledgeGraph } from '../graph/types'; -import { NODE_SCHEMA, EDGE_SCHEMA, NODE_TABLE_NAME, EDGE_TABLE_NAME } from './schema'; +import { NODE_SCHEMA, EDGE_SCHEMA, EMBEDDING_SCHEMA, NODE_TABLE_NAME, EDGE_TABLE_NAME } from './schema'; import { generateNodeCSV, generateEdgeCSV } from './csv-generator'; // Holds the reference to the dynamically loaded module @@ -34,8 +34,11 @@ export const initKuzu = async () => { // 3. Initialize WASM await kuzu.init(); - // 4. Create Database - db = new kuzu.Database(':memory:'); + // 4. Create Database with 512MB buffer pool + // Larger buffer needed for embedding storage (6K+ nodes ร— 384 floats) + // Constructor: Database(path, bufferPoolSize, maxNumThreads, enableCompression, readOnly) + const BUFFER_POOL_SIZE = 512 * 1024 * 1024; // 512MB + db = new kuzu.Database(':memory:', BUFFER_POOL_SIZE); conn = new kuzu.Connection(db); if (import.meta.env.DEV) console.log('โœ… KuzuDB WASM Initialized'); @@ -44,6 +47,7 @@ export const initKuzu = async () => { try { await conn.query(NODE_SCHEMA); await conn.query(EDGE_SCHEMA); + await conn.query(EMBEDDING_SCHEMA); if (import.meta.env.DEV) console.log('โœ… KuzuDB Schema Created'); } catch { // Schema might already exist, skip @@ -190,3 +194,160 @@ export const closeKuzu = async (): Promise => { } kuzu = null; }; + +/** + * Execute a prepared statement with parameters + * @param cypher - Cypher query with $param placeholders + * @param params - Object mapping param names to values + * @returns Query results + */ +export const executePrepared = async ( + cypher: string, + params: Record +): Promise => { + if (!conn) { + await initKuzu(); + } + + try { + // Note: conn.prepare is async in kuzu-wasm + const stmt = await conn.prepare(cypher); + if (!stmt.isSuccess()) { + const errMsg = await stmt.getErrorMessage(); + throw new Error(`Prepare failed: ${errMsg}`); + } + + const result = await conn.execute(stmt, params); + + // Collect all rows + const rows: any[] = []; + while (await result.hasNext()) { + const row = await result.getNext(); + rows.push(row); + } + + await stmt.close(); + return rows; + } catch (error) { + if (import.meta.env.DEV) console.error('Prepared query failed:', error); + throw error; + } +}; + +/** + * Execute a prepared statement with multiple parameter sets in small sub-batches + * Recreates statement every SUB_BATCH_SIZE executions to allow memory cleanup + * @param cypher - Cypher query with $param placeholders + * @param paramsList - Array of parameter objects to execute + */ +export const executeWithReusedStatement = async ( + cypher: string, + paramsList: Array> +): Promise => { + if (!conn) { + await initKuzu(); + } + + if (paramsList.length === 0) return; + + // Small sub-batch to allow memory cleanup between statement recreations + const SUB_BATCH_SIZE = 4; + + for (let i = 0; i < paramsList.length; i += SUB_BATCH_SIZE) { + const subBatch = paramsList.slice(i, i + SUB_BATCH_SIZE); + + // Create fresh statement for each sub-batch + const stmt = await conn.prepare(cypher); + if (!stmt.isSuccess()) { + const errMsg = await stmt.getErrorMessage(); + throw new Error(`Prepare failed: ${errMsg}`); + } + + try { + for (const params of subBatch) { + await conn.execute(stmt, params); + } + } finally { + await stmt.close(); + } + + // Small delay to allow garbage collection between sub-batches + if (i + SUB_BATCH_SIZE < paramsList.length) { + await new Promise(r => setTimeout(r, 0)); + } + } +}; + +/** + * Test if array parameters work with prepared statements + * This is a diagnostic function to check KuzuDB WASM capabilities + */ +export const testArrayParams = async (): Promise<{ success: boolean; error?: string }> => { + if (!conn) { + await initKuzu(); + } + + try { + // Test with a simple array parameter + const testEmbedding = new Array(384).fill(0).map((_, i) => i / 384); + + // First, get any node ID to test with + const nodeResult = await conn.query(`MATCH (n:${NODE_TABLE_NAME}) RETURN n.id AS id LIMIT 1`); + const nodeRow = await nodeResult.getNext(); + + if (!nodeRow) { + return { success: false, error: 'No nodes found to test with' }; + } + + const testNodeId = nodeRow.id ?? nodeRow[0]; + + if (import.meta.env.DEV) { + console.log('๐Ÿงช Testing array params with node:', testNodeId); + console.log('๐Ÿงช Embedding sample (first 5):', testEmbedding.slice(0, 5)); + } + + // Try using prepared statement with array param + // Note: conn.prepare is async in kuzu-wasm + const cypher = `MATCH (n:${NODE_TABLE_NAME} {id: $nodeId}) SET n.embedding = $embedding`; + const stmt = await conn.prepare(cypher); + + // In async API, isSuccess() returns boolean directly + if (!stmt.isSuccess()) { + const errMsg = await stmt.getErrorMessage(); + return { success: false, error: `Prepare failed: ${errMsg}` }; + } + + // Execute with array parameter + await conn.execute(stmt, { + nodeId: testNodeId, + embedding: testEmbedding, + }); + + await stmt.close(); + + // Verify it was stored + const verifyResult = await conn.query( + `MATCH (n:${NODE_TABLE_NAME} {id: '${testNodeId}'}) RETURN n.embedding AS emb` + ); + const verifyRow = await verifyResult.getNext(); + const storedEmb = verifyRow?.emb ?? verifyRow?.[0]; + + if (storedEmb && Array.isArray(storedEmb) && storedEmb.length === 384) { + if (import.meta.env.DEV) { + console.log('โœ… Array params WORK! Stored embedding length:', storedEmb.length); + } + return { success: true }; + } else { + return { + success: false, + error: `Embedding not stored correctly. Got: ${typeof storedEmb}, length: ${storedEmb?.length}` + }; + } + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error); + if (import.meta.env.DEV) { + console.error('โŒ Array params test failed:', errorMsg); + } + return { success: false, error: errorMsg }; + } +}; diff --git a/src/core/kuzu/schema.ts b/src/core/kuzu/schema.ts index adea1d2e19..b849136add 100644 --- a/src/core/kuzu/schema.ts +++ b/src/core/kuzu/schema.ts @@ -10,11 +10,12 @@ export const NODE_TABLE_NAME = 'CodeNode'; export const EDGE_TABLE_NAME = 'CodeRelation'; +export const EMBEDDING_TABLE_NAME = 'CodeEmbedding'; /** * Node table schema * Stores all code elements: Files, Functions, Classes, etc. - * embedding column stores 384-dimensional vectors for semantic search + * Note: Embeddings stored separately to avoid copy-on-write overhead */ export const NODE_SCHEMA = ` CREATE NODE TABLE ${NODE_TABLE_NAME} ( @@ -25,16 +26,27 @@ CREATE NODE TABLE ${NODE_TABLE_NAME} ( startLine INT64, endLine INT64, content STRING, - embedding FLOAT[384], PRIMARY KEY (id) )`; +/** + * Separate embedding table - lightweight structure for vector storage + * This avoids copy-on-write issues when storing embeddings + * (UPDATEing nodes with large content fields would copy entire node) + */ +export const EMBEDDING_SCHEMA = ` +CREATE NODE TABLE ${EMBEDDING_TABLE_NAME} ( + nodeId STRING, + embedding FLOAT[384], + PRIMARY KEY (nodeId) +)`; + /** * Create vector index for semantic search * Uses HNSW (Hierarchical Navigable Small World) algorithm with cosine similarity */ export const CREATE_VECTOR_INDEX_QUERY = ` -CALL CREATE_VECTOR_INDEX('${NODE_TABLE_NAME}', 'code_embedding_idx', 'embedding', metric := 'cosine') +CALL CREATE_VECTOR_INDEX('${EMBEDDING_TABLE_NAME}', 'code_embedding_idx', 'embedding', metric := 'cosine') `; /** @@ -50,5 +62,5 @@ CREATE REL TABLE ${EDGE_TABLE_NAME} ( /** * All schema creation queries in order */ -export const SCHEMA_QUERIES = [NODE_SCHEMA, EDGE_SCHEMA]; +export const SCHEMA_QUERIES = [NODE_SCHEMA, EDGE_SCHEMA, EMBEDDING_SCHEMA]; diff --git a/src/hooks/useAppState.tsx b/src/hooks/useAppState.tsx index 0bd761d72f..59570932ad 100644 --- a/src/hooks/useAppState.tsx +++ b/src/hooks/useAppState.tsx @@ -79,6 +79,9 @@ interface AppState { semanticSearch: (query: string, k?: number) => Promise; semanticSearchWithContext: (query: string, k?: number, hops?: number) => Promise; isEmbeddingReady: boolean; + + // Debug/test methods + testArrayParams: () => Promise<{ success: boolean; error?: string }>; } const AppStateContext = createContext(null); @@ -251,6 +254,12 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => { return api.semanticSearchWithContext(query, k, hops); }, []); + const testArrayParams = useCallback(async (): Promise<{ success: boolean; error?: string }> => { + const api = apiRef.current; + if (!api) return { success: false, error: 'Worker not initialized' }; + return api.testArrayParams(); + }, []); + const toggleLabelVisibility = useCallback((label: NodeLabel) => { setVisibleLabels(prev => { if (prev.includes(label)) { @@ -300,6 +309,8 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => { semanticSearch, semanticSearchWithContext, isEmbeddingReady: embeddingStatus === 'ready', + // Debug + testArrayParams, }; return ( diff --git a/src/workers/ingestion.worker.ts b/src/workers/ingestion.worker.ts index 4087b655cc..c281fbaff0 100644 --- a/src/workers/ingestion.worker.ts +++ b/src/workers/ingestion.worker.ts @@ -187,7 +187,7 @@ const workerApi = { onProgress(progress); }; - await runEmbeddingPipeline(kuzu.executeQuery, progressCallback); + await runEmbeddingPipeline(kuzu.executeQuery, kuzu.executeWithReusedStatement, progressCallback); }, /** @@ -266,6 +266,18 @@ const workerApi = { isEmbeddingComplete = false; embeddingProgress = null; }, + + /** + * Test if KuzuDB supports array parameters in prepared statements + * This is a diagnostic function + */ + async testArrayParams(): Promise<{ success: boolean; error?: string }> { + const kuzu = await getKuzuAdapter(); + if (!kuzu.isKuzuReady()) { + return { success: false, error: 'Database not ready' }; + } + return kuzu.testArrayParams(); + }, }; // Expose the worker API to the main thread