feat: add batch operations for embeddings and chunks in the database

Author: Helweg

benchmarks/run.ts

@@ -140,6 +140,15 @@ class Database {
     this.inner.upsertEmbedding(contentHash, embedding, chunkText, model);
   }
 
+  upsertEmbeddingsBatch(items: Array<{
+    contentHash: string;
+    embedding: Buffer;
+    chunkText: string;
+    model: string;
+  }>): void {
+    this.inner.upsertEmbeddingsBatch(items);
+  }
+
   upsertChunk(data: {
     chunkId: string;
     contentHash: string;
@@ -153,10 +162,27 @@ class Database {
     this.inner.upsertChunk(data);
   }
 
+  upsertChunksBatch(chunks: Array<{
+    chunkId: string;
+    contentHash: string;
+    filePath: string;
+    startLine: number;
+    endLine: number;
+    nodeType?: string;
+    name?: string;
+    language: string;
+  }>): void {
+    this.inner.upsertChunksBatch(chunks);
+  }
+
   addChunksToBranch(branch: string, chunkIds: string[]): void {
     this.inner.addChunksToBranch(branch, chunkIds);
   }
 
+  addChunksToBranchBatch(branch: string, chunkIds: string[]): void {
+    this.inner.addChunksToBranchBatch(branch, chunkIds);
+  }
+
   getBranchChunkIds(branch: string): string[] {
     return this.inner.getBranchChunkIds(branch);
   }
@@ -507,18 +533,19 @@ async function runBenchmarks(): Promise<void> {
     // ============================================
     console.log("\n=== Database Performance (SQLite) ===");
 
-    const dbPath = path.join(tempDir, "benchmark.db");
-    const db = new Database(dbPath);
-
     const chunkCounts = [100, 1000, 5000, 10000];
 
     for (const chunkCount of chunkCounts) {
+      const dbPath = path.join(tempDir, `benchmark-${chunkCount}.db`);
+      const db = new Database(dbPath);
+      const dbBatch = new Database(path.join(tempDir, `benchmark-batch-${chunkCount}.db`));
+
       const embedding = Buffer.from(
         new Float32Array(generateRandomEmbedding(dimensions)).buffer
       );
 
       benchmark(
-        `Insert ${chunkCount} embeddings`,
+        `Insert ${chunkCount} embeddings (sequential)`,
         () => {
           for (let i = 0; i < chunkCount; i++) {
             db.upsertEmbedding(`hash-${chunkCount}-${i}`, embedding, `text-${i}`, "test-model");
@@ -528,8 +555,24 @@ async function runBenchmarks(): Promise<void> {
         { embeddings: chunkCount }
       );
 
+      const embeddingBatchItems = Array.from({ length: chunkCount }, (_, i) => ({
+        contentHash: `hash-batch-${chunkCount}-${i}`,
+        embedding,
+        chunkText: `text-${i}`,
+        model: "test-model",
+      }));
+
       benchmark(
-        `Insert ${chunkCount} chunks`,
+        `Insert ${chunkCount} embeddings (batch)`,
+        () => {
+          dbBatch.upsertEmbeddingsBatch(embeddingBatchItems);
+        },
+        1,
+        { embeddings: chunkCount }
+      );
+
+      benchmark(
+        `Insert ${chunkCount} chunks (sequential)`,
         () => {
           for (let i = 0; i < chunkCount; i++) {
             db.upsertChunk({
@@ -548,19 +591,52 @@ async function runBenchmarks(): Promise<void> {
         { chunks: chunkCount }
       );
 
+      const chunkBatchItems = Array.from({ length: chunkCount }, (_, i) => ({
+        chunkId: `chunk-batch-${chunkCount}-${i}`,
+        contentHash: `hash-batch-${chunkCount}-${i}`,
+        filePath: `/file${i % 100}.ts`,
+        startLine: i * 10,
+        endLine: i * 10 + 10,
+        nodeType: "function",
+        name: `func${i}`,
+        language: "typescript",
+      }));
+
+      benchmark(
+        `Insert ${chunkCount} chunks (batch)`,
+        () => {
+          dbBatch.upsertChunksBatch(chunkBatchItems);
+        },
+        1,
+        { chunks: chunkCount }
+      );
+
       const chunkIds = Array.from(
         { length: chunkCount },
         (_, i) => `chunk-${chunkCount}-${i}`
       );
       benchmark(
-        `Add ${chunkCount} chunks to branch`,
+        `Add ${chunkCount} chunks to branch (sequential)`,
         () => {
           db.addChunksToBranch(`branch-${chunkCount}`, chunkIds);
         },
         1,
         { chunks: chunkCount }
       );
 
+      const chunkIdsBatch = Array.from(
+        { length: chunkCount },
+        (_, i) => `chunk-batch-${chunkCount}-${i}`
+      );
+      benchmark(
+        `Add ${chunkCount} chunks to branch (batch)`,
+        () => {
+          dbBatch.addChunksToBranchBatch(`branch-batch-${chunkCount}`, chunkIdsBatch);
+        },
+        1,
+        { chunks: chunkCount }
+      );
+
       benchmark(
         `Get branch chunk IDs (${chunkCount})`,
         () => {

native/src/db.rs

@@ -150,6 +150,35 @@ pub fn upsert_embedding(
     Ok(())
 }
 
+/// Batch insert or update embeddings within a single transaction
+pub fn upsert_embeddings_batch(
+    conn: &mut Connection,
+    embeddings: &[(String, Vec<u8>, String, String)],
+) -> DbResult<()> {
+    if embeddings.is_empty() {
+        return Ok(());
+    }
+
+    let tx = conn.transaction()?;
+    {
+        let mut stmt = tx.prepare(
+            r#"
+            INSERT INTO embeddings (content_hash, embedding, chunk_text, model, created_at)
+            VALUES (?, ?, ?, ?, strftime('%s', 'now'))
+            ON CONFLICT(content_hash) DO UPDATE SET
+                embedding = excluded.embedding,
+                model = excluded.model
+            "#,
+        )?;
+
+        for (content_hash, embedding, chunk_text, model) in embeddings {
+            stmt.execute(params![content_hash, embedding, chunk_text, model])?;
+        }
+    }
+    tx.commit()?;
+    Ok(())
+}
+
 /// Get multiple embeddings by content hashes
 #[allow(dead_code)]
 pub fn get_embeddings_batch(
@@ -250,6 +279,49 @@ pub fn upsert_chunk(
     Ok(())
 }
 
+/// Batch insert or update chunks within a single transaction
+pub fn upsert_chunks_batch(
+    conn: &mut Connection,
+    chunks: &[ChunkRow],
+) -> DbResult<()> {
+    if chunks.is_empty() {
+        return Ok(());
+    }
+
+    let tx = conn.transaction()?;
+    {
+        let mut stmt = tx.prepare(
+            r#"
+            INSERT INTO chunks (chunk_id, content_hash, file_path, start_line, end_line, node_type, name, language)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            ON CONFLICT(chunk_id) DO UPDATE SET
+                content_hash = excluded.content_hash,
+                file_path = excluded.file_path,
+                start_line = excluded.start_line,
+                end_line = excluded.end_line,
+                node_type = excluded.node_type,
+                name = excluded.name,
+                language = excluded.language
+            "#,
+        )?;
+
+        for chunk in chunks {
+            stmt.execute(params![
+                chunk.chunk_id,
+                chunk.content_hash,
+                chunk.file_path,
+                chunk.start_line,
+                chunk.end_line,
+                chunk.node_type,
+                chunk.name,
+                chunk.language
+            ])?;
+        }
+    }
+    tx.commit()?;
+    Ok(())
+}
+
 /// Get chunk by ID
 pub fn get_chunk(conn: &Connection, chunk_id: &str) -> DbResult<Option<ChunkRow>> {
     let result = conn
@@ -344,6 +416,30 @@ pub fn add_chunks_to_branch(conn: &Connection, branch: &str, chunk_ids: &[String
     Ok(())
 }
 
+/// Batch add chunks to a branch within a single transaction
+pub fn add_chunks_to_branch_batch(
+    conn: &mut Connection,
+    branch: &str,
+    chunk_ids: &[String],
+) -> DbResult<()> {
+    if chunk_ids.is_empty() {
+        return Ok(());
+    }
+
+    let tx = conn.transaction()?;
+    {
+        let mut stmt = tx.prepare(
+            "INSERT OR IGNORE INTO branch_chunks (branch, chunk_id) VALUES (?, ?)",
+        )?;
+
+        for chunk_id in chunk_ids {
+            stmt.execute(params![branch, chunk_id])?;
+        }
+    }
+    tx.commit()?;
+    Ok(())
+}
+
 /// Remove all chunks from a branch (for re-indexing)
 pub fn clear_branch(conn: &Connection, branch: &str) -> DbResult<usize> {
     let count = conn.execute(

native/src/lib.rs

@@ -263,6 +263,14 @@ pub struct BranchDelta {
     pub removed: Vec<String>,
 }
 
+#[napi(object)]
+pub struct EmbeddingBatchItem {
+    pub content_hash: String,
+    pub embedding: Buffer,
+    pub chunk_text: String,
+    pub model: String,
+}
+
 #[napi(object)]
 pub struct DatabaseStats {
     pub embedding_count: u32,
@@ -385,6 +393,51 @@ impl Database {
             .map_err(|e| Error::from_reason(e.to_string()))
     }
 
+    #[napi]
+    pub fn upsert_embeddings_batch(&self, items: Vec<EmbeddingBatchItem>) -> Result<()> {
+        let mut conn = self.conn.lock().map_err(|e| Error::from_reason(e.to_string()))?;
+        let batch: Vec<(String, Vec<u8>, String, String)> = items
+            .into_iter()
+            .map(|item| {
+                (
+                    item.content_hash,
+                    item.embedding.to_vec(),
+                    item.chunk_text,
+                    item.model,
+                )
+            })
+            .collect();
+        db::upsert_embeddings_batch(&mut conn, &batch)
+            .map_err(|e| Error::from_reason(e.to_string()))
+    }
+
+    #[napi]
+    pub fn upsert_chunks_batch(&self, chunks: Vec<ChunkData>) -> Result<()> {
+        let mut conn = self.conn.lock().map_err(|e| Error::from_reason(e.to_string()))?;
+        let batch: Vec<db::ChunkRow> = chunks
+            .into_iter()
+            .map(|c| db::ChunkRow {
+                chunk_id: c.chunk_id,
+                content_hash: c.content_hash,
+                file_path: c.file_path,
+                start_line: c.start_line,
+                end_line: c.end_line,
+                node_type: c.node_type,
+                name: c.name,
+                language: c.language,
+            })
+            .collect();
+        db::upsert_chunks_batch(&mut conn, &batch)
+            .map_err(|e| Error::from_reason(e.to_string()))
+    }
+
+    #[napi]
+    pub fn add_chunks_to_branch_batch(&self, branch: String, chunk_ids: Vec<String>) -> Result<()> {
+        let mut conn = self.conn.lock().map_err(|e| Error::from_reason(e.to_string()))?;
+        db::add_chunks_to_branch_batch(&mut conn, &branch, &chunk_ids)
+            .map_err(|e| Error::from_reason(e.to_string()))
+    }
+
     #[napi]
     pub fn clear_branch(&self, branch: String) -> Result<u32> {
         let conn = self.conn.lock().map_err(|e| Error::from_reason(e.to_string()))?;