Merge pull request #11 from VACLab/performance_improvement

improved cohort creation performance along with minor performance improvement in get concept hierarchy

Author: hyi

biasanalyzer/cohort.py

@@ -1,5 +1,6 @@
 from sqlalchemy.exc import SQLAlchemyError
 import duckdb
+import pandas as pd
 from datetime import datetime
 from pydantic import ValidationError
 from biasanalyzer.models import CohortDefinition, Cohort
@@ -110,14 +111,10 @@ def create_cohort(self, cohort_name: str, description: str, query_or_yaml_file:
             cohort_def_id = self.bias_db.create_cohort_definition(cohort_def)
 
             # Store cohort_definition and cohort data into BiasDatabase
-            for row in result:
-                cohort = Cohort(
-                    subject_id=int(row['person_id']),  # Assuming person_id column in the result
-                    cohort_definition_id=cohort_def_id,
-                    cohort_start_date=row['cohort_start_date'],
-                    cohort_end_date=row['cohort_end_date']
-                )
-                self.bias_db.create_cohort(cohort)
+            cohort_df = pd.DataFrame(result)
+            cohort_df['cohort_definition_id'] = cohort_def_id
+            cohort_df = cohort_df.rename(columns={"person_id": "subject_id"})
+            self.bias_db.create_cohort_in_bulk(cohort_df)
             print(f"Cohort {cohort_name} successfully created.")
             return CohortData(cohort_id=cohort_def_id, bias_db=self.bias_db, omop_db=self.omop_db)
         except duckdb.Error as e:

biasanalyzer/database.py

@@ -119,6 +119,15 @@ def create_cohort(self, cohort: Cohort):
             cohort.cohort_end_date
         ))
 
+    # Method to insert cohort data in bulk from a dataframe
+    def create_cohort_in_bulk(self, cohort_df: pd.DataFrame):
+        # make duckdb to treat cohort_df dataframe as a virtual table named "cohort_df"
+        self.conn.register("cohort_df", cohort_df)
+        self.conn.execute('''
+            INSERT INTO cohort (subject_id, cohort_definition_id, cohort_start_date, cohort_end_date)
+            SELECT subject_id, cohort_definition_id, cohort_start_date, cohort_end_date FROM cohort_df
+        ''')
+
     def get_cohort_definition(self, cohort_definition_id):
         results = self.conn.execute(f'''
         SELECT id, name, description, created_date, creation_info, created_by FROM cohort_definition 
@@ -417,12 +426,8 @@ def get_concept_hierarchy(self, concept_id: int):
 
         results = self.execute_query(query, params={"concept_id": concept_id})
 
-        # Collect all concept IDs involved in the hierarchy
-        concept_ids = set()
-        for row in results:
-            concept_ids.add(row['ancestor_concept_id'])
-            concept_ids.add(row['descendant_concept_id'])
-
+        # Collect all unique concept IDs involved in the hierarchy using set comprehension
+        concept_ids = {row['ancestor_concept_id'] for row in results} | {row['descendant_concept_id'] for row in results}
         # Fetch details of each concept
         concept_details = {}
         if concept_ids:
@@ -442,19 +447,17 @@ def get_concept_hierarchy(self, concept_id: int):
             ancestor_id = row['ancestor_concept_id']
             descendant_id = row['descendant_concept_id']
 
-            if ancestor_id not in hierarchy:
-                hierarchy[ancestor_id] = {"details": concept_details[ancestor_id], "children": []}
-            if descendant_id not in hierarchy:
-                hierarchy[descendant_id] = {"details": concept_details[descendant_id], "children": []}
-            # Link descendants to their ancestor node
-            hierarchy[ancestor_id]["children"].append(hierarchy[descendant_id])
-
-            if descendant_id not in reverse_hierarchy:
-                reverse_hierarchy[descendant_id] = {"details": concept_details[descendant_id], "parents": []}
-            if ancestor_id not in reverse_hierarchy:
-                reverse_hierarchy[ancestor_id] = {"details": concept_details[ancestor_id], "parents": []}
-            # Link ancestors to their descendant (child) node
-            reverse_hierarchy[descendant_id]["parents"].append(reverse_hierarchy[ancestor_id])
+            ancestor_entry = hierarchy.setdefault(
+                ancestor_id, {"details": concept_details[ancestor_id], "children": []})
+            descendant_entry = hierarchy.setdefault(
+                descendant_id, {"details": concept_details[descendant_id], "children": []})
+            ancestor_entry["children"].append(descendant_entry)
+
+            desc_entry_rev = reverse_hierarchy.setdefault(
+                descendant_id, {"details": concept_details[descendant_id], "parents": []})
+            ancestor_entry_rev = reverse_hierarchy.setdefault(
+                ancestor_id, {"details": concept_details[ancestor_id], "parents": []})
+            desc_entry_rev["parents"].append(ancestor_entry_rev)
 
         # Return the parent hierarchy and children hierarchy of the specified concept
         return reverse_hierarchy[concept_id], hierarchy[concept_id]