fix: prevent multiple filters to override spark plan

thehyve · Sep 17, 2024 · a358781 · a358781
1 parent 34aaaf5
commit a358781
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/src/gentropy/dataset/study_locus.py b/src/gentropy/dataset/study_locus.py
@@ -534,14 +534,16 @@ def filter_credible_set(
         Returns:
             StudyLocus: Filtered study-locus dataset.
         """
-        self.df = self._df.withColumn(
-            "locus",
-            f.filter(
-                f.col("locus"),
-                lambda tag: (tag[credible_interval.value]),
+        return StudyLocus(
+            _df=self._df.withColumn(
+                "locus",
+                f.filter(
+                    f.col("locus"),
+                    lambda tag: (tag[credible_interval.value]),
+                ),
             ),
+            _schema=self._schema,
         )
-        return self
 
     @staticmethod
     def filter_ld_set(ld_set: Column, r2_threshold: float) -> Column: