scriptoriumadmin/GitNexus
1
0
Fork 0
mirror of https://github.com/abhigyanpatwari/GitNexus.git synced 2026-10-29 06:32:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 7

embedings pipeline and storage done -> faced issue with buffer memory limit for kuzu, wrote about it in readme -> solved and working

Browse source
This commit is contained in:
abhigyanpatwari 2026-01-05 07:49:58 +05:30
parent c7437c3c81
commit a52f414138
7 changed files with 368 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

91
README.md
View file

@ -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<br/>id + name + content<br/>~2KB]
end
subgraph During["During UPDATE (COW)"]
N1_OLD[Old Version<br/>~2KB]
N1_NEW[New Version<br/>+ embedding<br/>~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<br/>id, name, content, embedding<br/>~3.5KB per UPDATE copy]
end
subgraph New["✅ New Design"]
CN2[CodeNode<br/>id, name, content]
CE[CodeEmbedding<br/>nodeId, embedding<br/>~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

52
src/components/EmbeddingStatus.tsx
View file

@ -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<string | null>(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 (
<button
onClick={handleStartEmbeddings}
className="flex items-center gap-2 px-3 py-1.5 bg-surface border border-border-subtle rounded-lg text-sm text-text-secondary hover:bg-hover hover:text-text-primary hover:border-accent/50 transition-all group"
title="Generate embeddings for semantic search"
>
<Brain className="w-4 h-4 text-node-interface group-hover:text-accent transition-colors" />
<span className="hidden sm:inline">Enable Semantic Search</span>
<Zap className="w-3 h-3 text-text-muted" />
</button>
<div className="flex items-center gap-2">
{/* Test button (dev only) */}
{import.meta.env.DEV && (
<button
onClick={handleTestArrayParams}
className="flex items-center gap-1 px-2 py-1.5 bg-surface border border-border-subtle rounded-lg text-xs text-text-muted hover:bg-hover hover:text-text-secondary transition-all"
title="Test if KuzuDB supports array params"
>
<FlaskConical className="w-3 h-3" />
{testResult || 'Test'}
</button>
)}
<button
onClick={handleStartEmbeddings}
className="flex items-center gap-2 px-3 py-1.5 bg-surface border border-border-subtle rounded-lg text-sm text-text-secondary hover:bg-hover hover:text-text-primary hover:border-accent/50 transition-all group"
title="Generate embeddings for semantic search"
>
<Brain className="w-4 h-4 text-node-interface group-hover:text-accent transition-colors" />
<span className="hidden sm:inline">Enable Semantic Search</span>
<Zap className="w-3 h-3 text-text-muted" />
</button>
</div>
);
}

71
src/core/embeddings/embedding-pipeline.ts
View file

@ -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<any[]>,
nodeId: string,
embedding: number[]
): Promise<void> => {
// 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<any[]>,
const batchInsertEmbeddings = async (
executeWithReusedStatement: (
cypher: string,
paramsList: Array<Record<string, any>>
) => Promise<void>,
updates: Array<{ id: string; embedding: number[] }>
): Promise<void> => {
// 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<any[]>
): Promise<void> => {
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<any[]>,
executeWithReusedStatement: (cypher: string, paramsList: Array<Record<string, any>>) => Promise<void>,
onProgress: EmbeddingProgressCallback,
config: Partial<EmbeddingConfig> = {}
): Promise<void> => {
@ -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,

167
src/core/kuzu/kuzu-adapter.ts
View file

@ -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<void> => {
}
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<string, any>
): Promise<any[]> => {
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<Record<string, any>>
): Promise<void> => {
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 };
}
};

20
src/core/kuzu/schema.ts
View file

@ -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];

11
src/hooks/useAppState.tsx
View file

@ -79,6 +79,9 @@ interface AppState {
semanticSearch: (query: string, k?: number) => Promise<SemanticSearchResult[]>;
semanticSearchWithContext: (query: string, k?: number, hops?: number) => Promise<any[]>;
isEmbeddingReady: boolean;
// Debug/test methods
testArrayParams: () => Promise<{ success: boolean; error?: string }>;
}
const AppStateContext = createContext<AppState | null>(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 (

14
src/workers/ingestion.worker.ts
View file

@ -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