Add vocabulary-grouped embedding counts and filtered concept-id enumeration

src/omop_emb/backends/base_backend.py

@@ -903,6 +903,35 @@ def get_embedding_count(
             )
         )
 
+    @require_registered_model
+    def get_embedding_count_by_vocabulary(
+        self,
+        *,
+        model_name: str,
+        metric_type: MetricType,
+        _model_record: EmbeddingModelRecord,
+    ) -> Mapping[str, int]:
+        """Return the number of stored embeddings, grouped by vocabulary_id.
+
+        Parameters
+        ----------
+        model_name : str
+            Canonical model name including tag.
+        metric_type : MetricType
+
+        Returns
+        -------
+        Mapping[str, int]
+            ``vocabulary_id`` to embedding count, for every vocabulary with at
+            least one stored embedding.
+        """
+        return self._get_embedding_count_by_vocabulary_impl(model_record=_model_record)
+
+    @abstractmethod
+    def _get_embedding_count_by_vocabulary_impl(
+        self, *, model_record: EmbeddingModelRecord
+    ) -> Mapping[str, int]: ...
+
     # ------------------------------------------------------------------
     # Validation helpers
     # ------------------------------------------------------------------

src/omop_emb/backends/pgvector/pg_backend.py

@@ -35,6 +35,7 @@
     q_all_concept_ids,
     q_concept_filter_metadata,
     q_concept_ids_matching_filter,
+    q_embedding_count_by_vocabulary,
     q_nearest_concept_ids,
     q_upsert_embeddings,
     q_create_extension_pgvector,
@@ -392,3 +393,11 @@ def _get_concept_ids_matching_filter_impl(
             setup_concept_filter_temps(session, concept_filter, "postgresql")
             rows = session.execute(q_concept_ids_matching_filter(table, concept_filter)).all()
         return {int(row[0]) for row in rows}
+
+    def _get_embedding_count_by_vocabulary_impl(
+        self, *, model_record: EmbeddingModelRecord
+    ) -> Mapping[str, int]:
+        table = self._table_cache[model_record.storage_identifier]
+        with self.emb_session_factory() as session:
+            rows = session.execute(q_embedding_count_by_vocabulary(table)).all()
+        return {row[0]: int(row[1]) for row in rows}

src/omop_emb/backends/pgvector/pg_sql.py

@@ -15,7 +15,7 @@
 from typing import List, Optional, Sequence, Union
 
 from numpy import ndarray
-from sqlalchemy import Engine, Integer, Select, inspect as sa_inspect, literal, select, text, TextClause
+from sqlalchemy import Engine, Integer, Select, func, inspect as sa_inspect, literal, select, text, TextClause
 from sqlalchemy.sql import cast, column, values
 from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.orm import mapped_column
@@ -141,6 +141,23 @@ def q_all_concept_ids(embedding_table: type[PGEmbeddingTable]) -> Select:
     return select(embedding_table.concept_id)
 
 
+def q_embedding_count_by_vocabulary(embedding_table: type[PGEmbeddingTable]) -> Select:
+    """Build a SELECT for stored embedding counts grouped by vocabulary_id.
+
+    Parameters
+    ----------
+    embedding_table : type[PGEmbeddingTable]
+
+    Returns
+    -------
+    Select
+        Columns: ``vocabulary_id`` (str), ``count`` (int).
+    """
+    return select(
+        embedding_table.vocabulary_id, func.count(embedding_table.concept_id)
+    ).group_by(embedding_table.vocabulary_id)
+
+
 def q_create_extension_pgvector() -> TextClause:
     """Return a SQL statement to create the pgvector extension if it doesn't exist."""
     return text("CREATE EXTENSION IF NOT EXISTS vector CASCADE;")

src/omop_emb/backends/sqlitevec/sqlitevec_backend.py

@@ -32,6 +32,7 @@
     dml_upsert_rows,
     query_all_concept_ids,
     query_concept_ids_matching_filter,
+    query_embedding_count_by_vocabulary,
     query_embeddings_by_ids,
     query_filter_metadata_by_ids,
     query_has_any,
@@ -286,3 +287,10 @@ def _get_concept_ids_matching_filter_impl(
                 table=table,
                 concept_filter=concept_filter,
             )
+
+    def _get_embedding_count_by_vocabulary_impl(
+        self, *, model_record: EmbeddingModelRecord
+    ) -> Mapping[str, int]:
+        table = self._table_cache[model_record.storage_identifier]
+        with self.emb_session_factory() as session:
+            return query_embedding_count_by_vocabulary(session=session, table=table)

src/omop_emb/backends/sqlitevec/sqlitevec_sql.py

@@ -293,6 +293,28 @@ def query_all_concept_ids(session: Session, table: Table) -> set[int]:
     return {int(row[0]) for row in rows}
 
 
+def query_embedding_count_by_vocabulary(session: Session, table: Table) -> dict[str, int]:
+    """Return stored embedding counts grouped by vocabulary_id.
+
+    Parameters
+    ----------
+    session : Session
+    table : Table
+
+    Returns
+    -------
+    dict[str, int]
+        ``vocabulary_id`` to embedding count, for every vocabulary with at
+        least one stored embedding.
+    """
+    stmt = (
+        select(table.c.vocabulary_id, func.count(table.c.concept_id))
+        .group_by(table.c.vocabulary_id)
+    )
+    rows = session.execute(stmt).all()
+    return {row[0]: int(row[1]) for row in rows}
+
+
 def query_embeddings_by_ids(
     session: Session,
     table: Table,

src/omop_emb/interface.py

@@ -222,6 +222,20 @@ def get_embedding_count(self) -> int:
             metric_type=self._metric_type,
         )
 
+    def get_embedding_count_by_vocabulary(self) -> Mapping[str, int]:
+        """Return stored embedding counts grouped by vocabulary_id.
+
+        Returns
+        -------
+        Mapping[str, int]
+            ``vocabulary_id`` to embedding count, for every vocabulary with at
+            least one stored embedding.
+        """
+        return self._backend.get_embedding_count_by_vocabulary(
+            model_name=self.canonical_model_name,
+            metric_type=self._metric_type,
+        )
+
     # ------------------------------------------------------------------
     # Search
     # ------------------------------------------------------------------
@@ -324,6 +338,29 @@ def get_embeddings_by_concept_ids(
             concept_ids=concept_ids,
         )
 
+    def get_indexed_concept_ids(
+        self,
+        concept_filter: Optional[EmbeddingConceptFilter] = None,
+    ) -> set[int]:
+        """Return every stored concept_id matching *concept_filter*.
+
+        Parameters
+        ----------
+        concept_filter : EmbeddingConceptFilter, optional
+            Filter constraints to evaluate (domain, vocabulary, standard,
+            concept ID allowlist). When omitted, every stored concept_id is
+            returned.
+
+        Returns
+        -------
+        set[int]
+        """
+        return self._backend.get_concept_ids_matching_filter(
+            model_name=self.canonical_model_name,
+            metric_type=self._metric_type,
+            concept_filter=concept_filter or EmbeddingConceptFilter(),
+        )
+
     # ------------------------------------------------------------------
     # Concepts without embedding (requires CDM)
     # ------------------------------------------------------------------

tests/shared_backend_tests.py

@@ -133,6 +133,14 @@ def test_upsert_count_matches(self, backend: EmbeddingBackend):
         )
         assert count == len(CONCEPT_RECORDS)
 
+    def test_get_embedding_count_by_vocabulary(self, backend: EmbeddingBackend):
+        self._upsert_all(backend)
+        counts = backend.get_embedding_count_by_vocabulary(
+            model_name=MODEL_NAME,
+            metric_type=MetricType.L2,
+        )
+        assert counts == {"SNOMED": 2, "RxNorm": 2}
+
     def test_upsert_is_idempotent(self, backend: EmbeddingBackend):
         self._upsert_all(backend)
         backend.upsert_embeddings(