diff --git a/src/omop_emb/backends/embedding_table.py b/src/omop_emb/backends/embedding_table.py index 654e67d..9a8591a 100644 --- a/src/omop_emb/backends/embedding_table.py +++ b/src/omop_emb/backends/embedding_table.py @@ -74,13 +74,15 @@ class ConceptEmbeddingRecord: Source vocabulary (e.g. ``'SNOMED'``, ``'RxNorm'``). is_standard : bool ``True`` if ``standard_concept`` is ``'S'`` or ``'C'``. + is_valid : bool + ``True`` if ``invalid_reason`` is not ``'D'`` or ``'U'``. """ concept_id: int domain_id: str vocabulary_id: str is_standard: bool - is_valid: bool = True + is_valid: bool class EmbeddingTableBase(DeclarativeBase): diff --git a/src/omop_emb/backends/pgvector/pg_backend.py b/src/omop_emb/backends/pgvector/pg_backend.py index bfd549c..48eede0 100644 --- a/src/omop_emb/backends/pgvector/pg_backend.py +++ b/src/omop_emb/backends/pgvector/pg_backend.py @@ -333,7 +333,10 @@ def _get_nearest_concepts_impl( NearestConceptMatch( concept_id=int(row.concept_id), similarity=float(similarity), + domain_id=row.domain_id, + vocabulary_id=row.vocabulary_id, is_standard=bool(row.is_standard), + is_active=bool(row.is_valid), ) ) diff --git a/src/omop_emb/backends/pgvector/pg_sql.py b/src/omop_emb/backends/pgvector/pg_sql.py index 12cbcf2..0bacfca 100644 --- a/src/omop_emb/backends/pgvector/pg_sql.py +++ b/src/omop_emb/backends/pgvector/pg_sql.py @@ -191,8 +191,14 @@ def q_nearest_concept_ids( Returns ------- Select - Columns: ``q_id`` (int), ``concept_id`` (int), ``is_standard`` (bool), ``distance`` (float). - Result shape is ``(Q*K, 4)`` before the caller re-groups by ``q_id``. + - ``q_id`` (int), + - ``concept_id`` (int), + - ``domain_id`` (str), + - ``vocabulary_id`` (str), + - ``is_standard`` (bool), + - ``is_valid`` (bool), + - ``distance`` (float). + Result shape is ``(Q*K, 7)`` before the caller re-groups by ``q_id``. Notes ----- @@ -213,7 +219,10 @@ def q_nearest_concept_ids( inner_stmt = ( select( embedding_table.concept_id, + embedding_table.domain_id, + embedding_table.vocabulary_id, embedding_table.is_standard, + embedding_table.is_valid, distance.label("distance"), ) .order_by(distance) @@ -233,7 +242,10 @@ def q_nearest_concept_ids( select( query_v.c.q_id, lateral_subq.c.concept_id, + lateral_subq.c.domain_id, + lateral_subq.c.vocabulary_id, lateral_subq.c.is_standard, + lateral_subq.c.is_valid, lateral_subq.c.distance, ) .select_from(query_v) diff --git a/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py b/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py index b96910b..dac33dd 100644 --- a/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py +++ b/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py @@ -231,11 +231,14 @@ def _get_nearest_concepts_impl( ) matches = tuple( NearestConceptMatch( - concept_id=concept_id, - similarity=get_similarity_from_distance(distance, metric_type), - is_standard=bool(is_standard), + concept_id=row.concept_id, + similarity=get_similarity_from_distance(row.distance, metric_type), + domain_id=row.domain_id, + vocabulary_id=row.vocabulary_id, + is_standard=bool(row.is_standard), + is_active=bool(row.is_valid), ) - for concept_id, distance, is_standard in rows + for row in rows ) results.append(matches) diff --git a/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py b/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py index d85f765..90b01af 100644 --- a/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py +++ b/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py @@ -8,16 +8,17 @@ import numpy as np from numpy import ndarray from sqlalchemy import ( - Column, - Engine, - Integer, - MetaData, - Table, - bindparam, - delete, - func, - insert, - select, + Column, + Engine, + Integer, + MetaData, + Row, + Table, + bindparam, + delete, + func, + insert, + select, text, LargeBinary ) @@ -208,7 +209,7 @@ def query_knn( metric_type: MetricType, k: int, concept_filter: Optional[EmbeddingConceptFilter] = None, -) -> list[tuple[int, float, int]]: +) -> Sequence[Row]: """Run a KNN query against a vec0 table using a per-query distance function. Parameters @@ -227,8 +228,17 @@ def query_knn( Returns ------- - list[tuple[int, float, int]] - Triples of ``(concept_id, distance, is_standard)`` ordered by distance ascending. + Sequence[Row] + - ``q_id`` (int), + - ``concept_id`` (int), + - ``domain_id`` (str), + - ``vocabulary_id`` (str), + - ``is_standard`` (bool), + - ``is_valid`` (bool), + - ``distance`` (float). + Orderer by ascending distance. Result shape is ``(≤k,)``. A row has fewer than + *k* entries only when fewer than *k* stored concepts exist that match + *concept_filter* (or exist at all). Raises ------ @@ -253,7 +263,18 @@ def query_knn( table.c[EMBEDDING_COLUMN_NAME], bindparam("q_emb", emb_blob) ).label("distance") - stmt = select(table.c.concept_id, distance, table.c.is_standard).order_by(distance).limit(k) + stmt = ( + select( + table.c.concept_id, + distance, + table.c.domain_id, + table.c.vocabulary_id, + table.c.is_standard, + table.c.is_valid, + ) + .order_by(distance) + .limit(k) + ) if concept_filter is not None: setup_concept_filter_temps(session, concept_filter, "sqlite") @@ -261,8 +282,7 @@ def query_knn( if concept_filter.limit is not None: stmt = stmt.limit(concept_filter.limit) - rows = session.execute(stmt).all() - return [(int(row[0]), float(row[1]), int(row[2])) for row in rows] + return session.execute(stmt).all() def query_concept_ids_matching_filter( diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index 6ff2579..0268aac 100644 --- a/src/omop_emb/interface.py +++ b/src/omop_emb/interface.py @@ -7,8 +7,12 @@ * Table identity is ``(model_name, provider_type)``: one row per model in the registry. ``metric_type`` is supplied by the caller at query time. * ``omop_cdm_engine`` is **optional** on the reader interface. When provided, - KNN results are enriched with concept names and flags from the CDM. - When absent, ``NearestConceptMatch.concept_name`` etc. are ``None``. + KNN results are enriched with ``concept_name`` from the CDM. When absent, + ``NearestConceptMatch.concept_name`` is ``None``. + ``domain_id``, ``vocabulary_id``, ``is_standard``, and ``is_active`` come + from the embedding table directly and are always populated regardless. + These attributes are necessary to be in the embedding table for filtering + without round-tripping to the CDM, and are populated from the CDM at ingestion time. * ``omop_cdm_engine`` is **required** for ingestion methods (``embed_and_upsert_concepts``) because concept metadata must be fetched from the CDM to populate the embedding table filter columns. @@ -105,8 +109,10 @@ class EmbeddingReaderInterface: Distance metric used for KNN queries and validated against the registry. omop_cdm_engine : Engine, optional Engine for the user's OMOP CDM. When provided, KNN results are - enriched with ``concept_name``, ``is_standard``, and ``is_active`` - from the CDM. When absent, those fields are ``None``. + enriched with ``concept_name`` from the CDM. When absent, + ``concept_name`` is ``None``. ``domain_id``, ``vocabulary_id``, + ``is_standard``, and ``is_active`` are populated directly from the + embedding table by the backend. model : str Model name in canonical form. canonical_model_name : str, optional @@ -262,7 +268,12 @@ def get_nearest_concepts( Returns ------- Tuple[Tuple[NearestConceptMatch, ...], ...] - Shape ``(Q, ≤k)``. Enrichment fields are ``None`` if no CDM engine. + Shape ``(Q, ≤k)``. A row has fewer than *k* entries only when fewer + than *k* stored concepts exist that match *concept_filter* (or + exist at all). ``domain_id``, ``vocabulary_id``, ``is_standard``, + and ``is_active`` are always populated from the embedding table. + ``concept_name`` is ``None`` if no CDM engine was provided to the + interface. """ effective_k = k or (concept_filter.limit if concept_filter else None) or self._k @@ -449,8 +460,10 @@ def _enrich( ) -> Tuple[Tuple[NearestConceptMatch, ...], ...]: """Enrich backend results with concept names from the CDM. - Only ``concept_name`` is populated here. ``is_standard`` is already - set by the backend from the embedding table filter columns. + Notes + ----- + Enriched concepts (if a CDM engine is provided) will have the following attributes: + - `concept_name` (str): The name of the concept from the CDM. """ if not self._cdm_engine: return raw diff --git a/src/omop_emb/storage/faiss/faiss_cache.py b/src/omop_emb/storage/faiss/faiss_cache.py index 7d43cdd..b073c12 100644 --- a/src/omop_emb/storage/faiss/faiss_cache.py +++ b/src/omop_emb/storage/faiss/faiss_cache.py @@ -39,7 +39,7 @@ from datetime import datetime, timezone from pathlib import Path from collections import OrderedDict -from typing import Optional, Tuple +from typing import Mapping, Optional, Tuple, cast import numpy as np from tqdm import tqdm @@ -425,14 +425,21 @@ def search( concept_filter : EmbeddingConceptFilter, optional Applied as a pre-filter over the live backend. backend : EmbeddingBackend, optional - Required when concept_filter is set and non-empty. Used to - resolve the filter to a concept-ID set via - :meth:`EmbeddingBackend.get_concept_ids_matching_filter`. + Used to resolve concept_filter to a concept-ID set (via + :meth:`EmbeddingBackend.get_concept_ids_matching_filter`) and, + regardless of concept_filter, to populate each result's + ``domain_id``, ``vocabulary_id``, ``is_standard``, and + ``is_active`` from the embedding table (via + :meth:`EmbeddingBackend.get_concept_filter_metadata`). FAISS + indices store only vectors and IDs, so these fields stay + ``None`` when backend is omitted. Required when + concept_filter is set and non-empty. Returns ------- tuple[tuple[NearestConceptMatch, ...], ...] - Shape ``(Q, <=k)``. CDM enrichment is the caller's responsibility. + Shape ``(Q, <=k)``. ``concept_name`` (the only CDM-sourced field) + is left ``None``. CDM enrichment is the caller's esponsibility. Notes ----- @@ -460,21 +467,38 @@ def search( distances, ids_matrix = index.search(query, k, params=params) + metadata_by_id: dict[int, Mapping[str, object]] = {} + if backend is not None: + unique_ids = {int(cid) for id_row in ids_matrix for cid in id_row if cid != -1} + if unique_ids: + metadata_by_id = backend.get_concept_filter_metadata( + model_name=self._model_name, + metric_type=metric_type, + concept_ids=tuple(unique_ids), + ) + results: list[tuple[NearestConceptMatch, ...]] = [] for dist_row, id_row in zip(distances, ids_matrix): - row_results = tuple( - NearestConceptMatch( - concept_id=int(cid), - similarity=float( - get_similarity_from_distance( - self._to_metric_dist(float(dist), metric_type), metric_type - ) - ), + row_matches = [] + for dist, cid in zip(dist_row, id_row): + if cid == -1: + continue + meta = metadata_by_id.get(int(cid), {}) + row_matches.append( + NearestConceptMatch( + concept_id=int(cid), + similarity=float( + get_similarity_from_distance( + self._to_metric_dist(float(dist), metric_type), metric_type + ) + ), + domain_id=cast(Optional[str], meta.get("domain_id")), + vocabulary_id=cast(Optional[str], meta.get("vocabulary_id")), + is_standard=cast(Optional[bool], meta.get("is_standard")), + is_active=cast(Optional[bool], meta.get("is_valid")), + ) ) - for dist, cid in zip(dist_row, id_row) - if cid != -1 - ) - results.append(row_results) + results.append(tuple(row_matches)) return tuple(results) # ------------------------------------------------------------------ diff --git a/src/omop_emb/utils/embedding_utils.py b/src/omop_emb/utils/embedding_utils.py index 06c1d67..aa0b125 100644 --- a/src/omop_emb/utils/embedding_utils.py +++ b/src/omop_emb/utils/embedding_utils.py @@ -122,8 +122,12 @@ class NearestConceptMatch: """Single nearest-neighbour result as returned to callers. ``concept_id`` and ``similarity`` are always populated by the backend. - The remaining fields are optionally enriched by the interface layer from - the OMOP CDM when an ``omop_cdm_engine`` is provided. + ``domain_id``, ``vocabulary_id``, ``is_standard``, and ``is_active`` are + also populated by the backend directly from the embedding table's filter + columns (see :class:`~omop_emb.backends.embedding_table.ConceptEmbeddingMixin`), + so they are available regardless of whether a CDM engine is configured. + ``concept_name`` is the only field enriched by the interface layer from + the OMOP CDM, and only when an ``omop_cdm_engine`` is provided. Attributes ---------- @@ -133,17 +137,27 @@ class NearestConceptMatch: Similarity score in ``[0.0, 1.0]``. Higher is more similar. concept_name : str, optional Human-readable concept name. ``None`` when no CDM engine is provided. + domain_id : str, optional + OMOP domain (e.g. ``'Condition'``, ``'Drug'``), taken from the + embedding table. ``None`` only if the backend could not resolve it. + vocabulary_id : str, optional + Source vocabulary (e.g. ``'SNOMED'``, ``'RxNorm'``), taken from the + embedding table. ``None`` only if the backend could not resolve it. is_standard : bool, optional - ``True`` if ``standard_concept`` is ``'S'`` or ``'C'``. ``None`` - when no CDM engine is provided. + ``True`` if ``standard_concept`` is ``'S'`` or ``'C'``, taken from + the embedding table. ``None`` only if the backend could not resolve + it. is_active : bool, optional - ``True`` if ``invalid_reason`` is not ``'D'`` or ``'U'``. ``None`` - when no CDM engine is provided. + ``True`` if ``invalid_reason`` is not ``'D'`` or ``'U'``, taken from + the embedding table. ``None`` only if the backend could not resolve + it. """ concept_id: int similarity: float concept_name: Optional[str] = None + domain_id: Optional[str] = None + vocabulary_id: Optional[str] = None is_standard: Optional[bool] = None is_active: Optional[bool] = None diff --git a/tests/conftest.py b/tests/conftest.py index fd4cd90..ec614ef 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -28,16 +28,16 @@ # This makes L2 and cosine tests fully deterministic. CONCEPT_RECORDS: tuple[ConceptEmbeddingRecord, ...] = ( ConceptEmbeddingRecord( - concept_id=1, domain_id="Condition", vocabulary_id="SNOMED", is_standard=True + concept_id=1, domain_id="Condition", vocabulary_id="SNOMED", is_standard=True, is_valid=True ), ConceptEmbeddingRecord( - concept_id=2, domain_id="Condition", vocabulary_id="SNOMED", is_standard=True + concept_id=2, domain_id="Condition", vocabulary_id="SNOMED", is_standard=True, is_valid=True ), ConceptEmbeddingRecord( - concept_id=3, domain_id="Drug", vocabulary_id="RxNorm", is_standard=True + concept_id=3, domain_id="Drug", vocabulary_id="RxNorm", is_standard=True, is_valid=True ), ConceptEmbeddingRecord( - concept_id=4, domain_id="Drug", vocabulary_id="RxNorm", is_standard=False + concept_id=4, domain_id="Drug", vocabulary_id="RxNorm", is_standard=False, is_valid=True ), ) diff --git a/tests/shared_backend_tests.py b/tests/shared_backend_tests.py index 1a2ce4b..ae12c35 100644 --- a/tests/shared_backend_tests.py +++ b/tests/shared_backend_tests.py @@ -194,6 +194,25 @@ def test_knn_returns_results(self, backend: EmbeddingBackend): assert len(results) == 1 assert len(results[0]) == 3 + def test_knn_returns_concept_filter_metadata(self, backend: EmbeddingBackend): + """domain_id/vocabulary_id/is_standard/is_active must come back on + every match, populated straight from the embedding table (no CDM + engine involved at this layer).""" + self._upsert_all(backend) + results = backend.get_nearest_concepts( + model_name=MODEL_NAME, + metric_type=MetricType.L2, + query_embeddings=QUERY_EMBEDDING, + k=len(CONCEPT_RECORDS), + ) + matches_by_id = {m.concept_id: m for m in results[0]} + for record in CONCEPT_RECORDS: + match = matches_by_id[record.concept_id] + assert match.domain_id == record.domain_id + assert match.vocabulary_id == record.vocabulary_id + assert match.is_standard == record.is_standard + assert match.is_active == record.is_valid + def test_knn_top1_is_closest(self, backend: EmbeddingBackend): # Query [-1] is closest to Diabetes [0] under L2 self._upsert_all(backend) diff --git a/tests/test_embedding_bundle.py b/tests/test_embedding_bundle.py index cd5cdf0..afdae92 100644 --- a/tests/test_embedding_bundle.py +++ b/tests/test_embedding_bundle.py @@ -48,7 +48,7 @@ def _populate( ) records = [ ConceptEmbeddingRecord( - concept_id=i, domain_id="Drug", vocabulary_id="RxNorm", is_standard=True + concept_id=i, domain_id="Drug", vocabulary_id="RxNorm", is_standard=True, is_valid=True ) for i in ids ] diff --git a/tests/test_faiss_cache.py b/tests/test_faiss_cache.py index 358be73..4beb0dd 100644 --- a/tests/test_faiss_cache.py +++ b/tests/test_faiss_cache.py @@ -42,7 +42,7 @@ _COSINE_RECORDS = [ ConceptEmbeddingRecord( - concept_id=i, domain_id="Test", vocabulary_id="Test", is_standard=True + concept_id=i, domain_id="Test", vocabulary_id="Test", is_standard=True, is_valid=True ) for i in _COSINE_IDS ] @@ -60,7 +60,7 @@ _L2_RECORDS = [ ConceptEmbeddingRecord( - concept_id=i, domain_id="Test", vocabulary_id="Test", is_standard=True + concept_id=i, domain_id="Test", vocabulary_id="Test", is_standard=True, is_valid=True ) for i in _L2_IDS ] @@ -457,7 +457,7 @@ def test_selective_filter_returns_full_top_k(self, tmp_path): records = [ ConceptEmbeddingRecord( concept_id=i, domain_id="Drug" if i == 5 else "Condition", - vocabulary_id="Test", is_standard=True, + vocabulary_id="Test", is_standard=True, is_valid=True, ) for i in range(1, 6) ] @@ -484,6 +484,32 @@ def test_selective_filter_returns_full_top_k(self, tmp_path): returned_ids = {r.concept_id for r in results[0]} assert returned_ids == {5} + def test_search_populates_concept_filter_metadata(self, tmp_path): + """When a backend is passed, matches must carry domain_id/vocabulary_id/ + is_standard/is_active resolved from the embedding table — FAISS itself + stores only vectors and IDs, so this depends on the metadata lookup.""" + backend = _make_svec_backend(_L2_DIM) + _populate_backend( + backend, dim=_L2_DIM, records=_L2_RECORDS, vecs=_L2_VECS, + metric_type=MetricType.L2, + ) + cache = _build_faiss(tmp_path, backend, MetricType.L2) + + results = cache.search( + _L2_QUERY, + k=len(_L2_RECORDS), + metric_type=MetricType.L2, + index_config=FlatIndexConfig(), + backend=backend, + ) + matches_by_id = {m.concept_id: m for m in results[0]} + for record in _L2_RECORDS: + match = matches_by_id[record.concept_id] + assert match.domain_id == record.domain_id + assert match.vocabulary_id == record.vocabulary_id + assert match.is_standard == record.is_standard + assert match.is_active == record.is_valid + def test_search_with_filter_requires_backend(self, tmp_path): backend = _make_svec_backend(_L2_DIM) _populate_backend( @@ -550,7 +576,7 @@ def test_is_fresh_false_after_later_upsert(self, tmp_path): metric_type=MetricType.L2, records=[ ConceptEmbeddingRecord( - concept_id=99, domain_id="Test", vocabulary_id="Test", is_standard=True + concept_id=99, domain_id="Test", vocabulary_id="Test", is_standard=True, is_valid=True ) ], embeddings=np.array([[5, 0]], dtype=np.float32), diff --git a/tests/test_sqlitevec.py b/tests/test_sqlitevec.py index 655c05b..9ce31ca 100644 --- a/tests/test_sqlitevec.py +++ b/tests/test_sqlitevec.py @@ -74,9 +74,14 @@ def test_flat_model_accepts_cosine_metric_at_query_time( domain_id="Condition", vocabulary_id="SNOMED", is_standard=True, + is_valid=True, ), ConceptEmbeddingRecord( - concept_id=3, domain_id="Drug", vocabulary_id="RxNorm", is_standard=True + concept_id=3, + domain_id="Drug", + vocabulary_id="RxNorm", + is_standard=True, + is_valid=True, ), ] nonzero_embeddings = np.array([[-10.0], [10.0]], dtype=np.float32)