sliftutils 1.4.86 → 1.4.87

package/index.d.ts

@@ -2393,7 +2393,7 @@ declare module "sliftutils/storage/proxydatabase/ivfEmbeddingDatabase" {
     };
     export declare function rebuildStructure(database: Database<IvfEmbeddingRoot>): void;
     export declare function searchEmbeddings(database: Database<IvfEmbeddingRoot>, query: StoredEmbedding, options: {
-        probeBudget: number;
+        probeCellCount: number;
         resultCount: number;
     }): SearchHit[] | undefined;
     export declare function lookupEmbeddings(database: Database<IvfEmbeddingRoot>, refs: string[]): Map<string, StoredEmbedding> | undefined;

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "sliftutils",
-    "version": "1.4.86",
+    "version": "1.4.87",
     "main": "index.js",
     "license": "MIT",
     "files": [

package/storage/proxydatabase/ivfDbCheck.ts

@@ -72,7 +72,7 @@ async function main() {
         const sampleCount = 100;
         for (let index = 0; index < sampleCount; index++) {
             const query = allInputs[Math.floor(index * run.size / sampleCount)];
-            const hits = searchEmbeddings(database, query.embedding, { probeBudget: 512, resultCount: 1 });
+            const hits = searchEmbeddings(database, query.embedding, { probeCellCount: 64, resultCount: 1 });
             if (hits && hits.length && hits[0].ref === query.ref) {
                 found++;
             }

package/storage/proxydatabase/ivfEmbeddingDatabase.d.ts

@@ -31,7 +31,7 @@ export type SearchHit = {
 };
 export declare function rebuildStructure(database: Database<IvfEmbeddingRoot>): void;
 export declare function searchEmbeddings(database: Database<IvfEmbeddingRoot>, query: StoredEmbedding, options: {
-    probeBudget: number;
+    probeCellCount: number;
     resultCount: number;
 }): SearchHit[] | undefined;
 export declare function lookupEmbeddings(database: Database<IvfEmbeddingRoot>, refs: string[]): Map<string, StoredEmbedding> | undefined;

package/storage/proxydatabase/ivfEmbeddingDatabase.ts

@@ -218,7 +218,8 @@ export function rebuildStructure(database: Database<IvfEmbeddingRoot>): void {
 export function searchEmbeddings(
     database: Database<IvfEmbeddingRoot>,
     query: StoredEmbedding,
-    options: { probeBudget: number; resultCount: number },
+    // NOTE: We do a progressive search, so probe cell count can be very high.
+    options: { probeCellCount: number; resultCount: number },
 ): SearchHit[] | undefined {
     const config = database.readData(root => root.config);
     if (!config) return undefined;
@@ -239,13 +240,36 @@ export function searchEmbeddings(
     const centroids = transactionRead(centroidStore(database));
     if (!centroids) return undefined;
     if (!centroids.size) return [];
-    const probeCellCount = Math.max(1, Math.ceil(options.probeBudget / config.cellTargetSize));
-    const probeCellIds = rankCellsByCloseness(query, centroids).slice(0, probeCellCount);
-    const stores = probeCellIds.map(cellId => database.readData(root => root.cells[cellId])).filter(isDefined);
+    const probeCellCount = Math.min(centroids.size, options.probeCellCount);
+
+    let rankedCellIds = rankCellsByCloseness(query, centroids);
+    let stores: TransactionSetStore<StoredEmbedding>[] = [];
+    // A somewhat strange search, trying to load more and more data. The reason we can't load everything at once is because the server has no way of knowing what we want first, and so it'll just wait until it can give us everything at once. By waiting until all of the earlier values load before even requesting the later values, it means that we will receive the earlier values sooner. Which will result in the search immediately returning results with very few accesses, And then the results improving in quality as time goes on, which is very desirable. 
+    {
+        let stopAtNextKeyPoint = false;
+        let keyPoints: number[] = [];
+        let currentProbeAdd = 4;
+        let currentTarget = currentProbeAdd;
+        while (currentTarget < probeCellCount) {
+            keyPoints.push(currentTarget);
+            currentTarget += currentProbeAdd;
+            currentProbeAdd *= 2;
+        }
+        for (let i = 0; i < probeCellCount; i++) {
+            if (keyPoints.includes(i) && stopAtNextKeyPoint) break;
+            let store = database.readData(root => root.cells[rankedCellIds[i]]);
+            if (!store) {
+                stopAtNextKeyPoint = true;
+                continue;
+            }
+            stores.push(store);
+        }
+    }
+
     for (const store of stores) {
         const members = replayTransactionStore<StoredEmbedding>(store);
-        for (const ref of members.keys()) {
-            hits.push({ ref, closeness: getCloseness(query, members.get(ref)!) });
+        for (const [ref, value] of members) {
+            hits.push({ ref, closeness: getCloseness(query, value) });
         }
     }
     hits.sort((left, right) => right.closeness - left.closeness);