fix: changes to get_major_population() function

opentargets · Mar 21, 2024 · b842bcf · b842bcf
1 parent ed6aa4e
commit b842bcf
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/src/gentropy/dataset/study_index.py b/src/gentropy/dataset/study_index.py
@@ -184,15 +184,14 @@ def has_summarystats(self: StudyIndex) -> Column:
         """
         return self.df.hasSumstats
 
-    def get_major_population(self: StudyIndex) -> DataFrame:
+    def get_major_population(self: StudyIndex) -> Column:
         """Extract major population from ldPopulationStructure rows with multiple ancestries.
 
         Returns:
-            DataFrame: Columns studyId and the extracted major population from ldPopulationStructure.
+            Column: Columns studyId and the extracted major population from ldPopulationStructure.
         """
-        return self.df.select(
-            "studyId",
-            f.array_max(f.col("ldPopulationStructure"))
+        return (
+            f.array_max(self.df.ldPopulationStructure)
             .getItem("ldPopulation")
-            .alias("majorPopulation"),
+            .alias("majorPopulation")
         )