sliftutils 1.4.87 → 1.4.88

package/index.d.ts

@@ -1153,6 +1153,7 @@ declare module "sliftutils/storage/BulkDatabase2/BulkDatabaseFormat" {
     export declare function buildFileBuffer(rows: Record<string, unknown>[], times: number[], targetBytes?: number): BuiltFile[];
     export declare function buildFileBufferRaw(rows: RawRow[], targetBytes?: number): BuiltFile[];
     export type BaseBulkDatabaseReader = {
+        name?: string;
         rowCount: number;
         totalBytes: number;
         minTime: number;
@@ -1194,6 +1195,7 @@ declare module "sliftutils/storage/BulkDatabase2/BulkDatabaseFormat" {
     export declare function loadBulkDatabase(config: {
         totalBytes: number;
         getRange: (start: number, end: number) => Promise<Buffer>;
+        name?: string;
     }): Promise<BaseBulkDatabaseReader>;
 
 }
@@ -2393,7 +2395,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: {
-        probeCellCount: number;
+        probeBudget: 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.87",
+    "version": "1.4.88",
     "main": "index.js",
     "license": "MIT",
     "files": [

package/storage/BulkDatabase2/BulkDatabaseFormat.d.ts

@@ -41,6 +41,7 @@ export interface BuiltFile {
 export declare function buildFileBuffer(rows: Record<string, unknown>[], times: number[], targetBytes?: number): BuiltFile[];
 export declare function buildFileBufferRaw(rows: RawRow[], targetBytes?: number): BuiltFile[];
 export type BaseBulkDatabaseReader = {
+    name?: string;
     rowCount: number;
     totalBytes: number;
     minTime: number;
@@ -82,4 +83,5 @@ export declare function loadBulkHeader(getRange: (start: number, end: number) =>
 export declare function loadBulkDatabase(config: {
     totalBytes: number;
     getRange: (start: number, end: number) => Promise<Buffer>;
+    name?: string;
 }): Promise<BaseBulkDatabaseReader>;

package/storage/BulkDatabase2/BulkDatabaseFormat.ts

@@ -410,6 +410,9 @@ export function buildFileBufferRaw(rows: RawRow[], targetBytes = TARGET_FILE_BYT
 }
 
 export type BaseBulkDatabaseReader = {
+    // Identifies the source this reader came from (the bulk file name, or "(streams)") so the join can
+    // name the offending file when one of its reads fails. Undefined for readers built without one.
+    name?: string;
     rowCount: number;
     totalBytes: number;
     // Write-time bounds of this reader's data (0 if unknown — old bulk files).
@@ -468,6 +471,7 @@ export async function loadBulkHeader(getRange: (start: number, end: number) => P
 export async function loadBulkDatabase(config: {
     totalBytes: number;
     getRange: (start: number, end: number) => Promise<Buffer>;
+    name?: string;
 }): Promise<BaseBulkDatabaseReader> {
     const headerLength = (await config.getRange(0, 4)).readUInt32LE(0);
     if (headerLength <= 0 || headerLength > config.totalBytes) {
@@ -502,6 +506,7 @@ export async function loadBulkDatabase(config: {
         : keys.map(() => header.maxTime || 0);
 
     return {
+        name: config.name,
         rowCount,
         totalBytes: config.totalBytes,
         minTime: header.minTime || 0,

package/storage/BulkDatabase2/LoadedIndex.ts

@@ -257,7 +257,7 @@ export async function loadFileReader(name: string, storage: FileStorage, f: Bulk
     const raw = await makeRawGetRange(storage, f.fileName);
     const fileId = nullJoin(name, f.fileName);
     const opened = await blockCache.open(fileId, raw.size, raw.rawGetRange);
-    const reader = await loadBulkDatabase({ totalBytes: opened.uncompressedSize, getRange: opened.getRange });
+    const reader = await loadBulkDatabase({ totalBytes: opened.uncompressedSize, getRange: opened.getRange, name: f.fileName });
     cache.set(f.fileName, reader);
     return reader;
 }
@@ -327,6 +327,19 @@ function pruneSubCaches(bulkFiles: BulkFileInfo[], streamFiles: StreamFileInfo[]
     for (const n of subCaches.stream.keys()) if (!liveStream.has(n)) subCaches.stream.delete(n);
 }
 
+// A corrupt/unreadable column in ONE underlying file must not break the whole joined read. When a
+// source's column read throws, we drop that source for that column (falling through to older readers,
+// exactly as if the file never set the column) and warn once per (file, column) so the corruption is
+// visible without flooding the log on every rebuild.
+const warnedCorruptReads = new Set<string>();
+function warnCorruptRead(db: BaseBulkDatabaseReader, column: string, e: unknown): void {
+    const fileName = db.name || "(unknown source)";
+    const dedupeKey = nullJoin(fileName, column);
+    if (warnedCorruptReads.has(dedupeKey)) return;
+    warnedCorruptReads.add(dedupeKey);
+    console.warn(`${red("corrupt read")}: file ${fileName} column ${JSON.stringify(column)} could not be read - skipping it (its data is dropped from results): ${(e as Error).message}`);
+}
+
 function joinBulkDatabases(databases: BaseBulkDatabaseReader[]): ResolvedReader {
     const deleteTime = new Map<string, number>();
     for (const db of databases) {
@@ -367,8 +380,13 @@ function joinBulkDatabases(databases: BaseBulkDatabaseReader[]): ResolvedReader
         async getColumn(column) {
             const perReader = await Promise.all(databases.map(async db => {
                 if (!db.columns.some(c => c.column === column)) return undefined;
-                const entries = await db.getColumn(column);
-                return new Map(entries.map(e => [e.key, { value: e.value, time: e.time }]));
+                try {
+                    const entries = await db.getColumn(column);
+                    return new Map(entries.map(e => [e.key, { value: e.value, time: e.time }]));
+                } catch (e) {
+                    warnCorruptRead(db, column, e);
+                    return undefined;
+                }
             }));
             return keys.map(key => {
                 let bestTime = -Infinity;
@@ -391,7 +409,13 @@ function joinBulkDatabases(databases: BaseBulkDatabaseReader[]): ResolvedReader
             let found = false;
             for (const db of databases) {
                 if (!db.columns.some(c => c.column === column)) continue;
-                const r = await db.getSingleField(key, column);
+                let r;
+                try {
+                    r = await db.getSingleField(key, column);
+                } catch (e) {
+                    warnCorruptRead(db, column, e);
+                    continue;
+                }
                 if (r === ABSENT) continue;
                 if (r.time > bestTime) { bestTime = r.time; bestVal = r.value; found = true; }
             }

package/storage/BulkDatabase2/streamLog.ts

@@ -106,6 +106,7 @@ export function streamReaderFromEntries(entries: StreamEntry[], totalBytes: numb
     }
     const keyIndexMap = buildKeyIndex(keys);
     let reader: BaseBulkDatabaseReader = {
+        name: "(streams)",
         totalBytes,
         rowCount: keys.length,
         minTime: times.size ? minTime : 0,

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, { probeCellCount: 64, resultCount: 1 });
+            const hits = searchEmbeddings(database, query.embedding, { probeBudget: 512, 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: {
-    probeCellCount: number;
+    probeBudget: number;
     resultCount: number;
 }): SearchHit[] | undefined;
 export declare function lookupEmbeddings(database: Database<IvfEmbeddingRoot>, refs: string[]): Map<string, StoredEmbedding> | undefined;

package/storage/proxydatabase/ivfEmbeddingDatabase.ts

@@ -2,7 +2,6 @@ import cborx from "cbor-x";
 import { Database, namespaceDatabase } from "./Database";
 import { TransactionSetStore, transactionRead, transactionMutate, transactionDelete, replayTransactionStore } from "./transactionSet";
 import { StoredEmbedding, EmbeddingFormat, getCloseness, embeddingToFloat32, releaseFloat32, encodeEmbedding, hashEmbedding } from "../embeddingFormats";
-import { isDefined } from "../../misc/types";
 
 export type IvfConfig = {
     model: string;
@@ -218,8 +217,7 @@ export function rebuildStructure(database: Database<IvfEmbeddingRoot>): void {
 export function searchEmbeddings(
     database: Database<IvfEmbeddingRoot>,
     query: StoredEmbedding,
-    // NOTE: We do a progressive search, so probe cell count can be very high.
-    options: { probeCellCount: number; resultCount: number },
+    options: { probeBudget: number; resultCount: number },
 ): SearchHit[] | undefined {
     const config = database.readData(root => root.config);
     if (!config) return undefined;
@@ -240,36 +238,14 @@ export function searchEmbeddings(
     const centroids = transactionRead(centroidStore(database));
     if (!centroids) return undefined;
     if (!centroids.size) return [];
-    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);
-        }
-    }
-
+    const probeCellCount = Math.max(1, Math.ceil(options.probeBudget / config.cellTargetSize));
+    const probeCellIds = rankCellsByCloseness(query, centroids).slice(0, probeCellCount);
+    const stores = database.readData(root => probeCellIds.map(cellId => root.cells[cellId]));
+    if (!stores) return undefined;
     for (const store of stores) {
         const members = replayTransactionStore<StoredEmbedding>(store);
-        for (const [ref, value] of members) {
-            hits.push({ ref, closeness: getCloseness(query, value) });
+        for (const ref of members.keys()) {
+            hits.push({ ref, closeness: getCloseness(query, members.get(ref)!) });
         }
     }
     hits.sort((left, right) => right.closeness - left.closeness);