perf: parallelize update task

Author: I-Info

packages/service/common/vectorStore/pg/index.ts

@@ -169,13 +169,13 @@ class PgClass {
     const pg = await connectPg();
     return pg.query<{ id: string }>(sql);
   }
-  async query<T extends QueryResultRow = any>(sql: string) {
+  async query<T extends QueryResultRow = any>(sql: string, warning = true) {
     const pg = await connectPg();
     const start = Date.now();
     return pg.query<T>(sql).then((res) => {
       const time = Date.now() - start;
 
-      if (time > 300) {
+      if (warning && time > 300) {
         addLog.warn(`pg query time: ${time}ms, sql: ${sql}`);
       }
 

projects/app/src/pages/api/admin/inithalfvec.ts

@@ -26,37 +26,101 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
       console.log('halfvector column added');
     }
 
-    let rowsUpdated;
-    let retryCount = 0;
-    do {
-      try {
-        rowsUpdated = await PgClient.query(`
-            WITH updated AS (
-              UPDATE ${DatasetVectorTableName}
-              SET halfvector = vector::halfvec(1536)
-              WHERE id IN (
-                  SELECT id
-                  FROM ${DatasetVectorTableName}
-                  WHERE halfvector IS NULL
-                  LIMIT 1000
-              )
-              RETURNING 1
-            )
-            SELECT count(*) FROM updated;
-          `);
-        console.log('rowsUpdated:', rowsUpdated.rows[0].count);
-      } catch (error) {
-        console.error('Error updating halfvector:', error);
-        retryCount++;
+    const maxIdResult = await PgClient.query(
+      `SELECT MAX(id) as max_id FROM ${DatasetVectorTableName}`
+    );
+    const maxId: number = maxIdResult.rows[0].max_id;
+
+    if (!maxId) {
+      console.warn('No data in the table: empty max_id');
+      jsonRes(res, { code: 500, error: 'No data in the table: empty max_id' });
+      return;
+    }
+
+    const batchSize = 25;
+    const numBatches = Math.ceil(maxId / batchSize);
+
+    const tasks: (() => Promise<void>)[] = [];
+    let totalRowsUpdated = 0;
+    let lastLoggedTime = Date.now();
+    let lastLoggedRows = 0;
+
+    const logUpdateSpeed = () => {
+      const currentTime = Date.now();
+      const timeElapsed = (currentTime - lastLoggedTime) / 1000; // seconds
+      const rowsUpdated = totalRowsUpdated - lastLoggedRows;
+      const speed = rowsUpdated / timeElapsed; // rows per second
+      console.log(`Update speed: ${speed.toFixed(2)} rows/s`);
+      lastLoggedTime = currentTime;
+      lastLoggedRows = totalRowsUpdated;
+    };
+
+    for (let i = 0; i < numBatches; i++) {
+      const startId = i * batchSize;
+      const endId = startId + batchSize;
+
+      const asyncUpdate = async () => {
+        let retryCount = 0;
+        do {
+          try {
+            const rowsUpdated = await PgClient.query(
+              `
+                UPDATE ${DatasetVectorTableName}
+                SET halfvector = vector::halfvec(1536)
+                WHERE id >= ${startId} AND id < ${endId} AND halfvector IS NULL;
+            `,
+              false
+            );
+            if (rowsUpdated?.rowCount) {
+              totalRowsUpdated += rowsUpdated.rowCount;
+              console.log(`Batch ${i + 1} - rowsUpdated: ${rowsUpdated.rowCount}`);
+            }
+            break;
+          } catch (error) {
+            console.error(`Error updating halfvector in batch ${i + 1}:`, error);
+            retryCount++;
+          }
+        } while (retryCount < 3);
+
+        if (retryCount >= 3) {
+          console.error(`Failed to update halfvector in batch ${i + 1} after 3 retries`);
+          Promise.reject(new Error('Failed to update halfvector in batch'));
+        }
+      };
+
+      tasks.push(asyncUpdate);
+    }
+
+    // randomize task list
+    tasks.sort(() => Math.random() - 0.5);
+
+    const executor = async (start: number, end: number) => {
+      console.log(`Executing tasks from ${start} to ${end}`);
+      for (let i = start; i < end; ++i) {
+        try {
+          await tasks[i]();
+        } catch (error) {
+          console.error(`Error updating halfvector in task ${i}`, error);
+        }
       }
-    } while (retryCount < 3 && rowsUpdated!.rows[0].count > 0);
-
-    if (retryCount >= 3) {
-      console.error('Failed to update halfvector after 3 retries');
-      return jsonRes(res, {
-        code: 500,
-        error: 'Failed to update halfvector after 3 retries'
-      });
+      console.log(`Tasks from ${start} to ${end} completed`);
+    };
+
+    const maxConcurrency = 20;
+    const taskSize = Math.ceil(tasks.length / maxConcurrency);
+    const promises = [];
+    for (let i = 0; i < maxConcurrency; ++i) {
+      const startId = i * taskSize;
+      const endId = Math.min((i + 1) * taskSize, tasks.length);
+      promises.push(executor(startId, endId));
+    }
+
+    const telemetryInterval = setInterval(logUpdateSpeed, 5000);
+
+    try {
+      await Promise.all(promises);
+    } finally {
+      clearInterval(telemetryInterval);
     }
 
     console.log('halfvector column updated');
@@ -70,6 +134,13 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
       COMMIT;
       `
     );
+    console.log('halfvector column set not null');
+
+    // 创建索引以提升查询性能
+    await PgClient.query(`
+      CREATE INDEX CONCURRENTLY IF NOT EXISTS halfvector_index ON ${DatasetVectorTableName} USING hnsw (halfvector halfvec_ip_ops) WITH (m = 32, ef_construction = 128);
+    `);
+    console.log('halfvector index created');
 
     // 后台释放空间，避免使用 VACUUM FULL 导致锁表。
     await PgClient.query(`VACUUM ${DatasetVectorTableName};`);