drop omop_table from the generated visit dataframe after using it

src/cehrbert_data/tools/ehrshot_to_omop.py

@@ -475,12 +475,14 @@ def generate_visit_id(data: DataFrame, time_interval: int = 12) -> DataFrame:
     )
 
     # We only allow the generated visit_ids associated with the visit_occurrence table
-    visit = data.select("patient_id", "visit_order", "omop_table").distinct().withColumn(
+    visit = data.where(
+        f.col("omop_table") == "visit_occurrence"
+    ).select("patient_id", "visit_order").distinct().withColumn(
         "new_visit_id",
         f.abs(
             f.hash(f.concat(f.col("patient_id").cast("string"), f.col("visit_order").cast("string")))
         ).cast("bigint")
-    ).where(f.col("omop_table") == "visit_occurrence")
+    )
 
     # Validate the uniqueness of visit_id
     visit.groupby("new_visit_id").count().select(f.assert_true(f.col("count") == 1))