From 667e5a9345603a920039df70882c5557a7907b83 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 10 Jul 2026 03:46:48 +0000 Subject: [PATCH] Implement catalogue status and querying by filter --- src/omop_emb/backends/base_backend.py | 29 +++++++++++++++ src/omop_emb/backends/pgvector/pg_backend.py | 9 +++++ src/omop_emb/backends/pgvector/pg_sql.py | 19 +++++++++- .../backends/sqlitevec/sqlitevec_backend.py | 8 ++++ .../backends/sqlitevec/sqlitevec_sql.py | 22 +++++++++++ src/omop_emb/interface.py | 37 +++++++++++++++++++ tests/shared_backend_tests.py | 8 ++++ 7 files changed, 131 insertions(+), 1 deletion(-) diff --git a/src/omop_emb/backends/base_backend.py b/src/omop_emb/backends/base_backend.py index d852d83..67c43cf 100644 --- a/src/omop_emb/backends/base_backend.py +++ b/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 # ------------------------------------------------------------------ diff --git a/src/omop_emb/backends/pgvector/pg_backend.py b/src/omop_emb/backends/pgvector/pg_backend.py index d801dcc..bfd549c 100644 --- a/src/omop_emb/backends/pgvector/pg_backend.py +++ b/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} diff --git a/src/omop_emb/backends/pgvector/pg_sql.py b/src/omop_emb/backends/pgvector/pg_sql.py index 80938f0..12cbcf2 100644 --- a/src/omop_emb/backends/pgvector/pg_sql.py +++ b/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;") diff --git a/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py b/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py index 9007c1a..b96910b 100644 --- a/src/omop_emb/backends/sqlitevec/sqlitevec_backend.py +++ b/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) diff --git a/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py b/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py index b2806d5..d85f765 100644 --- a/src/omop_emb/backends/sqlitevec/sqlitevec_sql.py +++ b/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, diff --git a/src/omop_emb/interface.py b/src/omop_emb/interface.py index 7b3c4e8..6ff2579 100644 --- a/src/omop_emb/interface.py +++ b/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) # ------------------------------------------------------------------ diff --git a/tests/shared_backend_tests.py b/tests/shared_backend_tests.py index a39cb65..1a2ce4b 100644 --- a/tests/shared_backend_tests.py +++ b/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(