fixed the bug in creating the new visit_id

knatarajan-lab · Jan 11, 2025 · 6215db4 · 6215db4
1 parent 9d0010f
commit 6215db4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/cehrbert_data/tools/ehrshot_to_omop.py b/src/cehrbert_data/tools/ehrshot_to_omop.py
@@ -667,13 +667,13 @@ def disconnect_visit_id(data: DataFrame, spark: SparkSession, cache_folder: str)
             .rowsBetween(Window.unboundedPreceding, Window.currentRow)
         )
     ).withColumn(
-        "row_number",
-        f.row_number().over(Window.orderBy(f.col("visit_id"), f.col("visit_partition")))
+        "visit_partition_rank",
+        f.dense_rank().over(Window.orderBy(f.col("visit_id"), f.col("visit_partition")))
     ).crossJoin(
         visit_records.select(f.max("visit_id").alias("max_visit_id"))
     ).withColumn(
         "new_visit_id",
-        f.col("max_visit_id") + f.col("row_number")
+        f.col("max_visit_id") + f.col("visit_partition_rank")
     ).drop(
         "max_visit_id", "row_number"
     )