File size: 5,217 Bytes
6c4ade3 4b63894 6c4ade3 e882a80 6c4ade3 710f03c 6c4ade3 710f03c 6c4ade3 e882a80 4b63894 e882a80 6c4ade3 e882a80 4b63894 e882a80 6c4ade3 e882a80 4b63894 6c4ade3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """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"):
# IVF indexes need this before .reconstruct(id) works at all --
# the default (Array) direct map assumes dense sequential ids,
# which ours (genome_id/gene_id) aren't; Hashtable handles
# arbitrary ids. Flat indexes have no such method/need.
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
|