Merge pull request #133 from SCAI-BIO/feat-sql-jsonl-import-export

feat: JSONL import/export functionality for SQL adapters

Author: mehmetcanay

datastew/process/jsonl_adapter.py

@@ -1,46 +1,38 @@
 import json
 import os
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Union
 
+import numpy as np
 import pandas as pd
 from tqdm import tqdm
 from weaviate.util import generate_uuid5
 
 from datastew.embedding import Vectorizer
-from datastew.repository import WeaviateRepository
+from datastew.repository import PostgreSQLRepository, SQLLiteRepository, WeaviateRepository
+from datastew.repository.base import BaseRepository
+from datastew.repository.model import Concept, Mapping, Terminology
 from datastew.repository.weaviate_schema import (
     concept_schema,
     mapping_schema_user_vectors,
     terminology_schema,
 )
 
 
-class WeaviateJsonlConverter(object):
-    """
-    Converts data to our JSONL format for Weaviate schema.
-    """
+class BaseJsonlConverter(ABC):
+    def __init__(self, dest_dir: str, buffer_size: int = 1000):
+        """Initialize the converter.
 
-    def __init__(
-        self,
-        dest_dir: str,
-        terminology_schema: dict = terminology_schema.schema,
-        concept_schema: dict = concept_schema.schema,
-        mapping_schema: dict = mapping_schema_user_vectors.schema,
-        buffer_size: int = 1000,
-    ):
+        :param dest_dir: Destination directory for the exported JSONL files.
+        :param buffer_size: Number of records to buffer before writing to disk, defaults to 1000
+        """
         self.dest_dir = dest_dir
-        self.terminology_schema = terminology_schema
-        self.concept_schema = concept_schema
-        self.mapping_schema = mapping_schema
         self._buffer = []
         self._buffer_size = buffer_size
         self._ensure_directories_exist()
 
     def _ensure_directories_exist(self):
-        """
-        Ensures the output directory exists.
-
-        :return: None
-        """
+        """Ensures the output directory exists."""
         os.makedirs(self.dest_dir, exist_ok=True)
 
     def _get_file_path(self, collection: str) -> str:
@@ -79,33 +71,187 @@ def _flush_to_file(self, file_path: str):
 
         with open(file_path, "a", encoding="utf-8") as file:
             for entry in self._buffer:
+                if isinstance(entry.get("embedding"), np.ndarray):
+                    entry["embedding"] = entry["embedding"].tolist()
                 file.write(json.dumps(entry) + "\n")
 
         self._buffer.clear()
 
-    def from_repository(self, repository: WeaviateRepository) -> None:
+    @abstractmethod
+    def from_repository(self, repository: BaseRepository):
+        pass
+
+    @abstractmethod
+    def from_ohdsi(self, src: str, vectorizer: Vectorizer = Vectorizer(), include_vectors: bool = True):
+        pass
+
+    @abstractmethod
+    def _object_to_dict(self, obj: Any) -> Dict[str, Any]:
+        pass
+
+
+class SQLJsonlConverter(BaseJsonlConverter):
+    def __init__(self, dest_dir: str, buffer_size: int = 1000):
+        super().__init__(dest_dir=dest_dir, buffer_size=buffer_size)
+
+    def from_repository(self, repository: Union[PostgreSQLRepository, SQLLiteRepository]):
+        """Export all records from a PostgreSQLRepository to JSONL files
+
+        :param repository: Active database repository instace.
+        """
+        session = repository.session
+
+        # Export Terminologies
+        terminology_file_path = self._get_file_path("terminology")
+        for t in tqdm(session.query(Terminology).all(), desc="Exporting Terminologies"):
+            terminology = self._object_to_dict(t)
+            self._write_to_jsonl(terminology_file_path, terminology)
+        self._flush_to_file(terminology_file_path)
+
+        # Export Concepts
+        concept_file_path = self._get_file_path("concept")
+        for c in tqdm(session.query(Concept).all(), desc="Exporting Concepts"):
+            concept = self._object_to_dict(c)
+            self._write_to_jsonl(concept_file_path, concept)
+        self._flush_to_file(concept_file_path)
+
+        # Export Mappings
+        mapping_file_path = self._get_file_path("mapping")
+        for m in tqdm(session.query(Mapping).all(), desc="Exporting Mappings"):
+            mapping = self._object_to_dict(m)
+            self._write_to_jsonl(mapping_file_path, mapping)
+        self._flush_to_file(mapping_file_path)
+
+    def from_ohdsi(self, src: str, vectorizer: Vectorizer = Vectorizer(), include_vectors: bool = True):
+        """
+        Converts data from OHDSI to SQL-compatible JSONL format.
+
+        :param src: Path to the OHDSI CONCEPT.csv file.
+        :param vectorizer: Vectorizer to use for text embeddings.
+        :param include_vectors: Whether to include vector data in mappings.
+        """
+        if not os.path.exists(src):
+            raise FileNotFoundError(f"OHDSI concept file '{src}' does not exist or is not a file.")
+
+        terminology_file_path = self._get_file_path("terminology")
+        concept_file_path = self._get_file_path("concept")
+        mapping_file_path = self._get_file_path("mapping")
+
+        # Write single OHDSI terminology entry
+        self._write_to_jsonl(terminology_file_path, {"id": "OHDSI", "name": "OHDSI"})
+        self._flush_to_file(terminology_file_path)
+
+        for chunk in tqdm(
+            pd.read_csv(
+                src,
+                delimiter="\t",
+                usecols=["concept_name", "concept_id"],
+                chunksize=10000,
+                dtype={"concept_id": str, "concept_name": str},
+            ),
+            desc="Processing OHDSI concepts",
+        ):
+
+            concepts = []
+            mappings = []
+
+            concept_names = chunk["concept_name"].astype(str).tolist()
+            concept_ids = chunk["concept_id"].astype(str).tolist()
+
+            if include_vectors:
+                embeddings = vectorizer.get_embeddings(concept_names)
+
+            for i in range(len(concept_names)):
+                concept_identifier = f"OHDSI:{concept_ids[i]}"
+                label = concept_names[i]
+
+                # Concept JSON
+                concepts.append(
+                    {
+                        "concept_identifier": concept_identifier,
+                        "pref_label": label,
+                        "terminology_id": "OHDSI",
+                    }
+                )
+
+                # Mapping JSON
+                mapping = {
+                    "concept_identifier": concept_identifier,
+                    "text": label,
+                }
+                if include_vectors:
+                    mapping["sentence_embedder"] = vectorizer.model_name
+                    mapping["embedding"] = embeddings[i]
+
+                mappings.append(mapping)
+
+            # Write results in batch
+            for concept_data in concepts:
+                self._write_to_jsonl(concept_file_path, concept_data)
+            self._flush_to_file(concept_file_path)
+            for mapping_data in mappings:
+                self._write_to_jsonl(mapping_file_path, mapping_data)
+            self._flush_to_file(mapping_file_path)
+
+    def _object_to_dict(self, obj: Union[Terminology, Concept, Mapping]) -> Dict[str, Any]:
+        if isinstance(obj, Terminology):
+            return {
+                "id": obj.id,
+                "name": obj.name,
+            }
+        elif isinstance(obj, Concept):
+            return {
+                "concept_identifier": obj.concept_identifier,
+                "pref_label": obj.pref_label,
+                "terminology_id": obj.terminology.id,
+            }
+        elif isinstance(obj, Mapping):
+            return {
+                "concept_identifier": obj.concept.concept_identifier,
+                "text": obj.text,
+                "embedding": obj.embedding,
+                "sentence_embedder": obj.sentence_embedder,
+            }
+        else:
+            raise TypeError(f"Unsupported object type: {type(obj)}")
+
+
+class WeaviateJsonlConverter(BaseJsonlConverter):
+    def __init__(
+        self,
+        dest_dir: str,
+        terminology_schema: dict = terminology_schema.schema,
+        concept_schema: dict = concept_schema.schema,
+        mapping_schema: dict = mapping_schema_user_vectors.schema,
+        buffer_size: int = 1000,
+    ):
+        super().__init__(dest_dir=dest_dir, buffer_size=buffer_size)
+        self.terminology_schema = terminology_schema
+        self.concept_schema = concept_schema
+        self.mapping_schema = mapping_schema
+
+    def from_repository(self, repository: WeaviateRepository):
         """
         Converts data from a WeaviateRepository to our JSONL format.
 
         :param repository: WeaviateRepository
-        :return: None
         """
         # Process terminology first
         terminology_file_path = self._get_file_path("terminology")
         for terminology in repository.get_iterator(self.terminology_schema["class"]):
-            self._write_to_jsonl(terminology_file_path, self._weaviate_object_to_dict(terminology))
+            self._write_to_jsonl(terminology_file_path, self._object_to_dict(terminology))
         self._flush_to_file(terminology_file_path)
 
         # Process concept next
         concept_file_path = self._get_file_path("concept")
         for concept in repository.get_iterator(self.concept_schema["class"]):
-            self._write_to_jsonl(concept_file_path, self._weaviate_object_to_dict(concept))
+            self._write_to_jsonl(concept_file_path, self._object_to_dict(concept))
         self._flush_to_file(concept_file_path)
 
         # Process mapping last
         mapping_file_path = self._get_file_path("mapping")
         for mapping in repository.get_iterator(self.mapping_schema["class"]):
-            self._write_to_jsonl(mapping_file_path, self._weaviate_object_to_dict(mapping))
+            self._write_to_jsonl(mapping_file_path, self._object_to_dict(mapping))
         self._flush_to_file(mapping_file_path)
 
     def from_ohdsi(self, src: str, vectorizer: Vectorizer = Vectorizer(), include_vectors: bool = True):
@@ -128,7 +274,7 @@ def from_ohdsi(self, src: str, vectorizer: Vectorizer = Vectorizer(), include_ve
         terminology_properties = {"name": "OHDSI"}
         terminology_id = generate_uuid5(terminology_properties)
         ohdsi_terminology = {
-            "class": self.terminology_schema["class"],
+            "class": "Terminology",
             "id": terminology_id,
             "properties": terminology_properties,
         }
@@ -206,22 +352,25 @@ def from_ohdsi(self, src: str, vectorizer: Vectorizer = Vectorizer(), include_ve
                 self._write_to_jsonl(mapping_file_path, mapping_data)
             self._flush_to_file(mapping_file_path)
 
-    @staticmethod
-    def _weaviate_object_to_dict(weaviate_object):
+    def _object_to_dict(self, obj) -> Dict[str, Any]:
+        """Conver a Weaviate object to a schema-compliant JSON dictionary.
 
-        if weaviate_object.references is not None:
+        :param obj: Weaviate object instance.
+        :return: Formatted dictionary containing class, id, properties, vector, and references.
+        """
+        if obj.references is not None:
             # FIXME: This is a hack to get the UUID of the referenced object. Replace as soon as weaviate devs offer an
             #  actual solution for this.
-            vals = [value.objects for _, value in weaviate_object.references.items()]
+            vals = [value.objects for _, value in obj.references.items()]
             uuid = [str(obj.uuid) for sublist in vals for obj in sublist][0]
-            references = {key: uuid for key, _ in weaviate_object.references.items()}
+            references = {key: uuid for key, _ in obj.references.items()}
         else:
             references = {}
 
         return {
-            "class": weaviate_object.collection,
-            "id": str(weaviate_object.uuid),
-            "properties": weaviate_object.properties,
-            "vector": weaviate_object.vector,
+            "class": obj.collection,
+            "id": str(obj.uuid),
+            "properties": obj.properties,
+            "vector": obj.vector,
             "references": references,
         }

datastew/repository/__init__.py

@@ -1,14 +1,6 @@
 from .model import Concept, Mapping, Terminology
+from .postgresql import PostgreSQLRepository
 from .sqllite import SQLLiteRepository
 from .weaviate import WeaviateRepository
-from .weaviate_schema import (concept_schema,
-                              mapping_schema_preconfigured_embeddings,
-                              mapping_schema_user_vectors, terminology_schema)
 
-__all__ = [
-    "Terminology",
-    "Concept",
-    "Mapping",
-    "SQLLiteRepository",
-    "WeaviateRepository",
-]
+__all__ = ["Terminology", "Concept", "Mapping", "SQLLiteRepository", "WeaviateRepository", "PostgreSQLRepository"]

datastew/repository/base.py

@@ -15,12 +15,12 @@ def __init__(self, vectorizer: Vectorizer = Vectorizer()):
         self.vectorizer = vectorizer
 
     @abstractmethod
-    def store(self, model_object_instance):
+    def store(self, model_object_instance: Union[Terminology, Concept, Mapping]):
         """Store a single model object instance."""
         pass
 
     @abstractmethod
-    def store_all(self, model_object_instances):
+    def store_all(self, model_object_instances: List[Union[Terminology, Concept, Mapping]]):
         """Store multiple model object instances."""
         pass
 
@@ -30,7 +30,7 @@ def get_concept(self, concept_id: str) -> Concept:
         pass
 
     @abstractmethod
-    def get_concepts(self) -> Page[Concept]:
+    def get_concepts(self, terminology_name: Optional[str] = None, offset: int = 0, limit: int = 100) -> Page[Concept]:
         """Retrieve all concepts from the database."""
         pass
 
@@ -95,6 +95,10 @@ def import_data_dictionary(self, data_dictionary: DataDictionarySource, terminol
             logger.exception("Failed to import data dictionary.")
             raise RuntimeError(f"Failed to import data dictionary source: {e}")
 
+    @abstractmethod
+    def import_from_jsonl(self, jsonl_path: str, object_type: str, chunk_size: int = 100):
+        pass
+
     def _parse_data_dictionary(
         self, data_dictionary: DataDictionarySource, terminology_name: str
     ) -> List[Union[Concept, Mapping, Terminology]]:

datastew/repository/model.py

@@ -73,7 +73,6 @@ class Concept(Base):
     pref_label = Column(String)
     terminology_id = Column(String, ForeignKey("terminology.id"))
     terminology = relationship("Terminology")
-    uuid = Column(String)
 
     def __init__(self, terminology: Terminology, pref_label: str, concept_identifier: str, id: Optional[str] = None):
         self.terminology = terminology