Include vocabulary and domain ID in NearestConceptMatch

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):

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),
                 )
             )
 

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)

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)
 

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,16 +263,26 @@ 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")
         stmt = apply_concept_filter_where(stmt, table.c, concept_filter)
         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(

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

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)
 
     # ------------------------------------------------------------------

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
 

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
     ),
 )