test: join based on variantId

opentargets · Feb 10, 2025 · 56c3c5c · 56c3c5c
1 parent 5f9b417
commit 56c3c5c
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/gentropy/dataset/l2g_features/intervals.py b/src/gentropy/dataset/l2g_features/intervals.py
@@ -35,6 +35,11 @@ def common_interval_feature_logic(
     """
     study_loci_exploded = (
         study_loci_to_annotate.df.withColumn("variantInLocus", f.explode_outer("locus"))
+        .select(
+            "studyLocusId",
+            f.col("variantInLocus.variantId").alias("variantId"),
+            f.col("variantInLocus.posteriorProbability").alias("posteriorProbability"),
+        )
         # Filter for PP > 0.001
         .filter(f.col("posteriorProbability") > 0.001)
         .select("studyLocusId", "variantId", "posteriorProbability")