| """Generic FAISS + DuckDB/Parquet nearest-neighbor store, mounted read-only |
| (hfmount). FAISS is mmap'd; the parquet table is never fully loaded -- |
| DuckDB queries it directly, pushing id filters down to row-group pruning. |
| |
| Two id modes, matching how an index was built: |
| id_column=None -- plain index, vector i <-> parquet row i (positional). |
| Table is small enough to load once (e.g. the ~10MB |
| annotation reference); searched ids ARE row positions. |
| id_column=<col> -- IndexIDMap2 built with add_with_ids(vecs, <col> values) |
| (e.g. genome_id/gene_id). Searched ids are real key |
| values; fetch_rows does a WHERE <col> IN (...) lookup |
| with no full-table load, for tables too big to load. |
| |
| One VectorStore is shared (lazy_singleton) across every concurrent request, |
| so every query goes through self._con.cursor() rather than self._con |
| directly -- a bare DuckDB connection isn't safe for concurrent queries from |
| multiple threads (confirmed empirically: concurrent .execute() calls on one |
| connection returned wrong/None results on ~20% of queries, silently, no |
| exception); a fresh cursor per call shares the connection but is safe. |
| """ |
| import duckdb |
| import faiss |
| import numpy as np |
|
|
| from utils.retry import retry |
|
|
|
|
| class VectorStore: |
| def __init__(self, index_path, parquet_path, id_column=None, top_k=16, preload=False): |
| """preload=True reads the whole index into process RAM up front |
| (one slower sequential read at construction) instead of mmap'ing it |
| (lazy, on-demand network fetch per touched page on every query) -- |
| worth it for an index small enough to fit comfortably in RAM, to |
| make every later query immune to the mount's latency/flakiness.""" |
| self.parquet_path = parquet_path |
| self.id_column = id_column |
| self.top_k = top_k |
| flags = 0 if preload else faiss.IO_FLAG_MMAP |
| self.index = retry(lambda: faiss.read_index(index_path, flags), label=f"read_index {index_path}") |
| if hasattr(self.index, "set_direct_map_type"): |
| |
| |
| |
| |
| self.index.set_direct_map_type(faiss.DirectMap.Hashtable) |
| self._normalize = (self.index.metric_type == faiss.METRIC_INNER_PRODUCT) |
| self._con = duckdb.connect() |
|
|
| if id_column is None: |
| self.table = retry( |
| lambda: self._con.cursor().execute("SELECT * FROM read_parquet(?)", [parquet_path]).df(), |
| label=f"load {parquet_path}") |
| n_rows = len(self.table) |
| else: |
| self.table = None |
| n_rows = retry( |
| lambda: self._con.cursor().execute("SELECT count(*) FROM read_parquet(?)", [parquet_path]).fetchone()[0], |
| label=f"count {parquet_path}") |
|
|
| if n_rows != self.index.ntotal: |
| raise RuntimeError( |
| f"row-count mismatch: {parquet_path} has {n_rows} rows, " |
| f"{index_path} has {self.index.ntotal} vectors." |
| ) |
|
|
| def search(self, vectors, k=None): |
| """vectors: [N, D] array-like. Returns (scores, ids), each [N, k] -- |
| ids are row positions (id_column=None) or key values (id_column set); |
| -1 pads rows with fewer than k hits.""" |
| vecs = np.asarray(vectors, dtype=np.float32) |
| if self._normalize: |
| norms = np.linalg.norm(vecs, axis=1, keepdims=True) |
| vecs = vecs / np.clip(norms, 1e-12, None) |
| return self.index.search(vecs, k or self.top_k) |
|
|
| def fetch_rows(self, ids): |
| """ids: iterable of row positions or key values (matching id_column |
| mode), -1 entries ignored. Returns {id: row_dict}.""" |
| ids = sorted({int(i) for i in ids if i >= 0}) |
| if not ids: |
| return {} |
| if self.id_column is None: |
| return {i: self.table.iloc[i].to_dict() for i in ids} |
| placeholders = ",".join("?" * len(ids)) |
| query = f"SELECT * FROM read_parquet(?) WHERE {self.id_column} IN ({placeholders})" |
|
|
| def _run(): |
| cur = self._con.cursor() |
| cur.execute(query, [self.parquet_path] + ids) |
| return cur.fetchall(), [d[0] for d in cur.description] |
|
|
| rows, cols = retry(_run, label=f"fetch_rows {self.parquet_path}") |
| return {row[cols.index(self.id_column)]: dict(zip(cols, row)) for row in rows} |
|
|
| def search_and_fetch(self, vectors, k=None): |
| """One search + one batched fetch_rows for the whole query batch. |
| Returns a list of per-query hit lists: [{**row, "score": float}, ...].""" |
| scores, ids = self.search(vectors, k) |
| rows = self.fetch_rows(ids.ravel()) |
| results = [] |
| for s_row, i_row in zip(scores, ids): |
| results.append([ |
| {**rows[int(i)], "score": float(s)} |
| for s, i in zip(s_row, i_row) if i >= 0 and int(i) in rows |
| ]) |
| return results |
|
|
