feat: only new keys for annotations

Author: bpblanken

v03_pipeline/lib/misc/clickhouse.py

@@ -81,9 +81,9 @@ def key_field(self):
     @property
     def join_condition(self):
         return (
-            'src.variantId = dst.variantId'
+            'assumeNotNull(src.variantId) = dst.variantId'
             if self == ClickHouseTable.KEY_LOOKUP
-            else 'toUInt32(src.key) = dst.key'
+            else 'assumeNotNull(toUInt32(src.key)) = dst.key'
         )
 
     @property
@@ -92,11 +92,32 @@ def select_fields(self):
 
     @property
     def insert(self) -> Callable:
-        return (
-            functools.partial(direct_insert, clickhouse_table=self)
-            if self != ClickHouseTable.ENTRIES
-            else functools.partial(atomic_entries_insert, _clickhouse_table=self)
-        )
+        return {
+            ClickHouseTable.ANNOTATIONS_DISK: functools.partial(
+                direct_insert_new_keys,
+                clickhouse_table=self,
+            ),
+            ClickHouseTable.ANNOTATIONS_MEMORY: functools.partial(
+                direct_insert_new_keys,
+                clickhouse_table=self,
+            ),
+            ClickHouseTable.CLINVAR: functools.partial(
+                direct_insert_all_keys,
+                clickhouse_table=self,
+            ),
+            ClickHouseTable.ENTRIES: functools.partial(
+                atomic_entries_insert,
+                _clickhouse_table=self,
+            ),
+            ClickHouseTable.KEY_LOOKUP: functools.partial(
+                direct_insert_all_keys,
+                _clickhouse_table=self,
+            ),
+            ClickHouseTable.TRANSCRIPTS: functools.partial(
+                direct_insert_all_keys,
+                _clickhouse_table=self,
+            ),
+        }[self]
 
 
 class ClickHouseDictionary(StrEnum):
@@ -414,7 +435,7 @@ def exchange_entity(
 
 
 @retry()
-def direct_insert(
+def direct_insert_new_keys(
     clickhouse_table: ClickHouseTable,
     reference_genome: ReferenceGenome,
     dataset_type: DatasetType,
@@ -456,6 +477,30 @@ def direct_insert(
     drop_staging_db()
 
 
+@retry()
+def direct_insert_all_keys(
+    clickhouse_table: ClickHouseTable,
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    run_id: str,
+    **_,
+) -> None:
+    table_name_builder = TableNameBuilder(
+        reference_genome,
+        dataset_type,
+        run_id,
+    )
+    dst_table = table_name_builder.dst_table(clickhouse_table)
+    src_table = table_name_builder.src_table(clickhouse_table)
+    logged_query(
+        f"""
+        INSERT INTO {dst_table}
+        SELECT {clickhouse_table.select_fields}
+        FROM {src_table}
+        """,
+    )
+
+
 @retry()
 def atomic_entries_insert(
     _clickhouse_table: ClickHouseTable,

v03_pipeline/lib/misc/clickhouse_test.py

@@ -15,7 +15,8 @@
     create_staging_non_table_entities,
     create_staging_tables,
     delete_existing_families_from_staging_entries,
-    direct_insert,
+    direct_insert_all_keys,
+    direct_insert_new_keys,
     exchange_entity,
     get_clickhouse_client,
     insert_new_entries,
@@ -234,7 +235,7 @@ def test_table_name_builder(self):
                 "gcs('https://storage.googleapis.com/mock_bucket/v3.1/GRCh38/SNV_INDEL/runs/manual__2025-05-07T17-20-59.702114+00-00/new_entries.parquet/*.parquet', '', '', 'Parquet')",
             )
 
-    def test_direct_insert(self):
+    def test_direct_insert_all_keys(self):
         client = get_clickhouse_client()
         df = pd.DataFrame({'key': [1, 2, 3, 4], 'transcripts': ['a', 'b', 'c', 'd']})
         table = pa.Table.from_pandas(df)
@@ -260,7 +261,7 @@ def test_direct_insert(self):
             f'INSERT INTO {Env.CLICKHOUSE_DATABASE}.`GRCh38/SNV_INDEL/transcripts` VALUES',
             [(1, 'a'), (10, 'b'), (7, 'c')],
         )
-        direct_insert(
+        direct_insert_all_keys(
             ClickHouseTable.TRANSCRIPTS,
             ReferenceGenome.GRCh38,
             DatasetType.SNV_INDEL,
@@ -275,7 +276,7 @@ def test_direct_insert(self):
         )
 
         # ensure multiple calls are idempotent
-        direct_insert(
+        direct_insert_all_keys(
             ClickHouseTable.TRANSCRIPTS,
             ReferenceGenome.GRCh38,
             DatasetType.SNV_INDEL,
@@ -289,7 +290,7 @@ def test_direct_insert(self):
             [(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd'), (7, 'c'), (10, 'b')],
         )
 
-    def test_direct_insert_key_lookup(self):
+    def test_direct_insert_key_lookup_new_keys(self):
         client = get_clickhouse_client()
         df = pd.DataFrame(
             {
@@ -325,7 +326,7 @@ def test_direct_insert_key_lookup(self):
             f'INSERT INTO {Env.CLICKHOUSE_DATABASE}.`GRCh38/SNV_INDEL/key_lookup` VALUES',
             [('1-123-A-C', 1), ('2-234-C-T', 2), ('M-345-C-G', 3)],
         )
-        direct_insert(
+        direct_insert_new_keys(
             ClickHouseTable.KEY_LOOKUP,
             ReferenceGenome.GRCh38,
             DatasetType.SNV_INDEL,