cehrbert-data 0.0.7__tar.gz → 0.0.8__tar.gz

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cehrbert_data
-Version: 0.0.7
+Version: 0.0.8
 Summary: The Spark ETL tools for generating the CEHR-BERT and CEHR-GPT pre-training and finetuning data
 Author-email: Chao Pang <chaopang229@gmail.com>, Xinzhuo Jiang <xj2193@cumc.columbia.edu>, Krishna Kalluri <kk3326@cumc.columbia.edu>, Nishanth Parameshwar Pavinkurve <np2689@cumc.columbia.edu>, Karthik Natarajan <kn2174@cumc.columbia.edu>
 License: MIT License

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/apps/generate_included_concept_list.py

@@ -24,8 +24,8 @@ from pyspark.sql import SparkSession
 from pyspark.sql import functions as F
 
 from cehrbert_data.config.output_names import QUALIFIED_CONCEPT_LIST_PATH
-from cehrbert_data.const.common import MEASUREMENT
-from cehrbert_data.utils.spark_utils import join_domain_tables, preprocess_domain_table
+from cehrbert_data.const.common import CONCEPT
+from cehrbert_data.utils.spark_utils import extract_events_by_domain, preprocess_domain_table
 
 DOMAIN_TABLE_LIST = ["condition_occurrence", "procedure_occurrence", "drug_exposure"]
 
@@ -55,32 +55,38 @@ def main(
     """
     spark = SparkSession.builder.appName("Generate concept list").getOrCreate()
 
-    domain_tables = []
     # Exclude measurement from domain_table_list if exists because we need to process measurement
     # in a different way
     domain_table_list = ehr_table_list if ehr_table_list else DOMAIN_TABLE_LIST
+    concept = preprocess_domain_table(spark, input_folder, CONCEPT)
+    patient_ehr_events = None
     for domain_table_name in domain_table_list:
-        if domain_table_name != MEASUREMENT:
-            domain_tables.append(
-                preprocess_domain_table(
-                    spark,
-                    input_folder,
-                    domain_table_name,
-                    with_drug_rollup=with_drug_rollup,
-                )
-            )
-
-    # Union all domain table records
-    patient_events = join_domain_tables(domain_tables)
+        domain_table = preprocess_domain_table(
+            spark=spark,
+            input_folder=input_folder,
+            domain_table_name=domain_table_name,
+            with_drug_rollup=with_drug_rollup
+        )
+        ehr_events = extract_events_by_domain(
+            domain_table,
+            spark=spark,
+            concept=concept,
+            aggregate_by_hour=False,
+            refresh=False,
+            persistence_folder=input_folder
+        )
+        if patient_ehr_events is None:
+            patient_ehr_events = ehr_events
+        else:
+            patient_ehr_events = patient_ehr_events.unionByName(ehr_events)
 
     # Filter out concepts that are linked to less than 100 patients
     qualified_concepts = (
-        patient_events.where("visit_occurrence_id IS NOT NULL")
+        patient_ehr_events.where("visit_occurrence_id IS NOT NULL")
         .groupBy("standard_concept_id")
         .agg(F.countDistinct("person_id").alias("freq"))
         .where(F.col("freq") >= min_num_of_patients)
     )
-
     qualified_concepts.write.mode("overwrite").parquet(os.path.join(output_folder, QUALIFIED_CONCEPT_LIST_PATH))
 
 

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/apps/generate_training_data.py

@@ -12,15 +12,14 @@ from cehrbert_data.const.common import (
     PERSON,
     VISIT_OCCURRENCE,
     DEATH,
-    MEASUREMENT
+    CONCEPT
 )
 from cehrbert_data.decorators import AttType
 from cehrbert_data.utils.spark_utils import (
     create_sequence_data,
     create_sequence_data_with_att,
-    join_domain_tables,
     preprocess_domain_table,
-    get_measurement_table,
+    extract_events_by_domain,
     validate_table_names,
     construct_artificial_visits
 )
@@ -43,7 +42,6 @@ def main(
         include_death: bool,
         att_type: AttType,
         inpatient_att_type: AttType,
-        include_sequence_information_content: bool = False,
         exclude_demographic: bool = False,
         use_age_group: bool = False,
         with_drug_rollup: bool = True,
@@ -52,6 +50,7 @@ def main(
         refresh_measurement: bool = False,
         aggregate_by_hour: bool = True,
         should_construct_artificial_visits: bool = False,
+        duplicate_records: bool = False
 ):
     spark = SparkSession.builder.appName("Generate CEHR-BERT Training Data").getOrCreate()
 
@@ -77,19 +76,30 @@ def main(
         f"refresh_measurement: {refresh_measurement}\n"
         f"aggregate_by_hour: {aggregate_by_hour}\n"
         f"should_construct_artificial_visits: {should_construct_artificial_visits}\n"
+        f"duplicate_records: {duplicate_records}\n"
     )
 
-    domain_tables = []
+    concept = preprocess_domain_table(spark, input_folder, CONCEPT)
+    patient_ehr_events = None
     for domain_table_name in domain_table_list:
-        if domain_table_name != MEASUREMENT:
-            domain_tables.append(
-                preprocess_domain_table(
-                    spark,
-                    input_folder,
-                    domain_table_name,
-                    with_drug_rollup=with_drug_rollup,
-                )
-            )
+        domain_table = preprocess_domain_table(
+            spark=spark,
+            input_folder=input_folder,
+            domain_table_name=domain_table_name,
+            with_drug_rollup=with_drug_rollup
+        )
+        ehr_events = extract_events_by_domain(
+            domain_table,
+            spark=spark,
+            concept=concept,
+            aggregate_by_hour=aggregate_by_hour,
+            refresh=refresh_measurement,
+            persistence_folder=input_folder
+        )
+        if patient_ehr_events is None:
+            patient_ehr_events = ehr_events
+        else:
+            patient_ehr_events = patient_ehr_events.unionByName(ehr_events)
 
     visit_occurrence = preprocess_domain_table(spark, input_folder, VISIT_OCCURRENCE)
     visit_occurrence = visit_occurrence.select(
@@ -119,35 +129,17 @@ def main(
 
     death = preprocess_domain_table(spark, input_folder, DEATH) if include_death else None
 
-    patient_events = join_domain_tables(domain_tables)
-
-    if include_concept_list and patient_events:
-        column_names = patient_events.schema.fieldNames()
+    if include_concept_list and patient_ehr_events:
         # Filter out concepts
-        qualified_concepts = preprocess_domain_table(spark, input_folder, "qualified_concept_list").select(
-            "standard_concept_id"
-        )
-
-        patient_events = patient_events.join(qualified_concepts, "standard_concept_id").select(column_names)
-
-    # Process the measurement table if exists
-    if MEASUREMENT in domain_table_list:
-        processed_measurement = get_measurement_table(
-            spark,
-            input_folder,
-            refresh=refresh_measurement,
-            aggregate_by_hour=aggregate_by_hour,
+        qualified_concepts = preprocess_domain_table(spark, input_folder, "qualified_concept_list")
+        patient_ehr_events = patient_ehr_events.join(
+            qualified_concepts.select("standard_concept_id"), "standard_concept_id"
         )
-        if patient_events:
-            # Union all measurement records together with other domain records
-            patient_events = patient_events.unionByName(processed_measurement)
-        else:
-            patient_events = processed_measurement
 
-    patient_events = (
-        patient_events.join(visit_occurrence_person, "visit_occurrence_id")
+    patient_ehr_events = (
+        patient_ehr_events.join(visit_occurrence_person, "visit_occurrence_id")
         .select(
-            [patient_events[fieldName] for fieldName in patient_events.schema.fieldNames()]
+            [patient_ehr_events[fieldName] for fieldName in patient_ehr_events.schema.fieldNames()]
             + ["visit_concept_id", "age"]
         )
         .withColumn("cohort_member_id", F.col("person_id"))
@@ -156,25 +148,35 @@ def main(
     # Apply the age security measure
     # We only keep the patient records, whose corresponding age is less than 90
     if apply_age_filter:
-        patient_events = patient_events.where(F.col("age") < 90)
+        patient_ehr_events = patient_ehr_events.where(F.col("age") < 90)
 
     if not continue_from_events:
-        patient_events.write.mode("overwrite").parquet(os.path.join(output_folder, "all_patient_events"))
-
-    patient_events = spark.read.parquet(os.path.join(output_folder, "all_patient_events"))
+        patient_ehr_events.write.mode("overwrite").parquet(os.path.join(output_folder, "all_patient_events"))
 
+    patient_ehr_events = spark.read.parquet(os.path.join(output_folder, "all_patient_events"))
     if should_construct_artificial_visits:
         # Construct artificial visits or re-link the visits for the problem list events
-        patient_events, visit_occurrence_person = construct_artificial_visits(
-            patient_events,
+        patient_ehr_events, visit_occurrence_person = construct_artificial_visits(
+            patient_ehr_events,
             visit_occurrence_person,
             spark=spark,
             persistence_folder=output_folder,
+            duplicate_records=duplicate_records
         )
+        # Update age if some of the ehr_records have been re-associated with the new visits
+        patient_ehr_events = patient_ehr_events.join(
+            person.select("person_id", "birth_datetime"),
+            "person_id",
+        ).join(
+            visit_occurrence_person.select("visit_occurrence_id", "visit_start_date"), "visit_occurrence_id"
+        ).withColumn(
+            "age",
+            F.ceil(F.months_between(F.col("visit_start_date"), F.col("birth_datetime")) / F.lit(12))
+        ).drop("visit_start_date", "birth_datetime")
 
     if is_new_patient_representation:
-        sequence_data = create_sequence_data_with_att(
-            patient_events,
+        patient_sequence_data = create_sequence_data_with_att(
+            patient_ehr_events,
             visit_occurrence_person,
             date_filter=date_filter,
             include_visit_type=include_visit_type,
@@ -190,8 +192,8 @@ def main(
             persistence_folder=output_folder,
         )
     else:
-        sequence_data = create_sequence_data(
-            patient_events,
+        patient_sequence_data = create_sequence_data(
+            patient_ehr_events,
             date_filter=date_filter,
             include_visit_type=include_visit_type,
             classic_bert_seq=is_classic_bert,
@@ -215,47 +217,33 @@ def main(
             )
             .distinct()
         )
-        sequence_data = sequence_data.join(visit_occurrence, "person_id")
-
-    if include_sequence_information_content:
-        concept_df = patient_events.select("person_id", F.col("standard_concept_id").alias("concept_id"))
-        concept_freq = (
-            concept_df.groupBy("concept_id")
-            .count()
-            .withColumn("prob", F.col("count") / F.sum("count").over(Window.partitionBy()))
-            .withColumn("ic", -F.log("prob"))
-        )
-
-        patient_ic_df = concept_df.join(concept_freq, "concept_id").groupby("person_id").agg(F.mean("ic").alias("ic"))
-
-        sequence_data = sequence_data.join(patient_ic_df, "person_id")
+        patient_sequence_data = patient_sequence_data.join(visit_occurrence, "person_id")
 
     patient_splits_folder = os.path.join(input_folder, "patient_splits")
     if os.path.exists(patient_splits_folder):
         patient_splits = spark.read.parquet(patient_splits_folder)
         temp_folder = os.path.join(output_folder, "patient_sequence", "temp")
-        sequence_data.join(
+        patient_sequence_data.join(
             patient_splits.select("person_id", "split"),
             "person_id"
         ).write.mode("overwrite").parquet(
             temp_folder
         )
-        sequence_data = spark.read.parquet(temp_folder)
-        sequence_data.where('split="train"').write.mode("overwrite").parquet(
+        patient_sequence_data = spark.read.parquet(temp_folder)
+        patient_sequence_data.where('split="train"').write.mode("overwrite").parquet(
             os.path.join(output_folder, "patient_sequence/train")
         )
-        sequence_data.where('split="test"').write.mode("overwrite").parquet(
+        patient_sequence_data.where('split="test"').write.mode("overwrite").parquet(
             os.path.join(output_folder, "patient_sequence/test")
         )
         shutil.rmtree(temp_folder)
     else:
-        sequence_data.write.mode("overwrite").parquet(os.path.join(output_folder, "patient_sequence"))
+        patient_sequence_data.write.mode("overwrite").parquet(os.path.join(output_folder, "patient_sequence"))
 
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Arguments for generate training data for Bert")
+def create_argparser():
+    parser = argparse.ArgumentParser(description="Arguments for generate training data for CEHR-Bert/GPT")
     parser.add_argument(
-        "-i",
         "--input_folder",
         dest="input_folder",
         action="store",
@@ -263,7 +251,6 @@ if __name__ == "__main__":
         required=True,
     )
     parser.add_argument(
-        "-o",
         "--output_folder",
         dest="output_folder",
         action="store",
@@ -271,7 +258,6 @@ if __name__ == "__main__":
         required=True,
     )
     parser.add_argument(
-        "-tc",
         "--domain_table_list",
         dest="domain_table_list",
         nargs="+",
@@ -369,8 +355,17 @@ if __name__ == "__main__":
         help="Indicate whether we should construct artificial visits for "
              "the problem list records that could occur years ahead",
     )
+    parser.add_argument(
+        "--duplicate_records",
+        dest="duplicate_records",
+        action="store_true",
+        help="Indicate whether we want to duplicate the problem list records when constructing artificial visits"
+    )
+    return parser
+
 
-    ARGS = parser.parse_args()
+if __name__ == "__main__":
+    ARGS = create_argparser().parse_args()
 
     # Enable logging
     add_console_logging()
@@ -399,4 +394,5 @@ if __name__ == "__main__":
         refresh_measurement=ARGS.refresh_measurement,
         aggregate_by_hour=ARGS.aggregate_by_hour,
         should_construct_artificial_visits=ARGS.should_construct_artificial_visits,
+        duplicate_records=ARGS.duplicate_records
     )

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/cohorts/spark_app_base.py

@@ -13,16 +13,17 @@ from pyspark.sql.window import Window
 from cehrbert_data.decorators import AttType
 from cehrbert_data.const.common import VISIT_OCCURRENCE
 from cehrbert_data.utils.spark_utils import (
-    build_ancestry_table_for,
     create_concept_frequency_data,
-    create_hierarchical_sequence_data,
     create_sequence_data,
     create_sequence_data_with_att,
     extract_ehr_records,
-    get_descendant_concept_ids,
     preprocess_domain_table,
     construct_artificial_visits
 )
+from cehrbert_data.utils.vocab_utils import (
+    build_ancestry_table_for,
+    get_descendant_concept_ids
+)
 from cehrbert_data.utils.logging_utils import add_console_logging
 
 from ..cohorts.query_builder import ENTRY_COHORT, NEGATIVE_COHORT, QueryBuilder
@@ -300,7 +301,6 @@ class NestedCohortBuilder:
             aggregate_by_hour: bool = True,
             is_new_patient_representation: bool = False,
             gpt_patient_sequence: bool = False,
-            is_hierarchical_bert: bool = False,
             classic_bert_seq: bool = False,
             is_first_time_outcome: bool = False,
             is_questionable_outcome_existed: bool = False,
@@ -318,6 +318,7 @@ class NestedCohortBuilder:
             meds_format: bool = False,
             cache_events: bool = False,
             should_construct_artificial_visits: bool = False,
+            duplicate_records: bool = False,
     ):
         self._cohort_name = cohort_name
         self._input_folder = input_folder
@@ -342,7 +343,6 @@ class NestedCohortBuilder:
         self._is_drug_roll_up_concept = is_drug_roll_up_concept
         self._is_new_patient_representation = is_new_patient_representation
         self._gpt_patient_sequence = gpt_patient_sequence
-        self._is_hierarchical_bert = is_hierarchical_bert
         self._is_first_time_outcome = is_first_time_outcome
         self._is_remove_index_prediction_starts = is_remove_index_prediction_starts
         self._is_questionable_outcome_existed = is_questionable_outcome_existed
@@ -365,6 +365,7 @@ class NestedCohortBuilder:
         self._meds_format = meds_format
         self._cache_events = cache_events
         self._should_construct_artificial_visits = should_construct_artificial_visits
+        self._duplicate_records = duplicate_records
 
         self.get_logger().info(
             f"cohort_name: {cohort_name}\n"
@@ -386,7 +387,6 @@ class NestedCohortBuilder:
             f"is_drug_roll_up_concept: {is_drug_roll_up_concept}\n"
             f"is_new_patient_representation: {is_new_patient_representation}\n"
             f"gpt_patient_sequence: {gpt_patient_sequence}\n"
-            f"is_hierarchical_bert: {is_hierarchical_bert}\n"
             f"is_first_time_outcome: {is_first_time_outcome}\n"
             f"is_questionable_outcome_existed: {is_questionable_outcome_existed}\n"
             f"is_remove_index_prediction_starts: {is_remove_index_prediction_starts}\n"
@@ -406,6 +406,7 @@ class NestedCohortBuilder:
             f"meds_format: {meds_format}\n"
             f"cache_events: {cache_events}\n"
             f"should_construct_artificial_visits: {should_construct_artificial_visits}\n"
+            f"duplicate_records: {duplicate_records}\n"
         )
 
         self.spark = SparkSession.builder.appName(f"Generate {self._cohort_name}").getOrCreate()
@@ -637,12 +638,36 @@ class NestedCohortBuilder:
             )
 
         if self._should_construct_artificial_visits:
+            person = self._dependency_dict[PERSON]
+            birthdate_udf = F.coalesce(
+                "birth_datetime",
+                F.concat("year_of_birth", F.lit("-01-01")).cast("timestamp"),
+            )
+            patient_demographic = person.select(
+                "person_id",
+                birthdate_udf.alias("birth_datetime"),
+            )
             ehr_records, visit_occurrence_with_artificial_visits = construct_artificial_visits(
                 ehr_records,
                 self._dependency_dict[VISIT_OCCURRENCE],
                 spark=self.spark if self._cache_events else None,
                 persistence_folder=self._output_data_folder if self._cache_events else None,
+                duplicate_records=self._duplicate_records
             )
+
+            # Update age if some of the ehr_records have been re-associated with the new visits
+            ehr_records = ehr_records.join(
+                patient_demographic,
+                "person_id"
+            ).join(
+                visit_occurrence_with_artificial_visits.select(
+                    "visit_occurrence_id", "visit_start_date"
+                ), "visit_occurrence_id"
+            ).withColumn(
+                "age",
+                F.ceil(F.months_between(F.col("visit_start_date"), F.col("birth_datetime")) / F.lit(12))
+            ).drop("visit_start_date", "birth_datetime")
+
             # Refresh the dependency
             self._dependency_dict[VISIT_OCCURRENCE] = visit_occurrence_with_artificial_visits
 
@@ -672,8 +697,9 @@ class NestedCohortBuilder:
             else:
                 if self._is_observation_window_unbounded:
                     record_window_filter = (
-                        F.col("ehr.datetime")
-                        <= F.expr(f"cohort.index_date - INTERVAL {self._hold_off_window} DAYS + INTERVAL 0.1 SECOND")
+                            F.col("ehr.datetime")
+                            <= F.expr(
+                        f"cohort.index_date - INTERVAL {self._hold_off_window} DAYS + INTERVAL 0.1 SECOND")
                     )
                 else:
                     record_window_filter = F.col("ehr.datetime").between(
@@ -692,14 +718,6 @@ class NestedCohortBuilder:
             [F.col("ehr." + field_name) for field_name in ehr_records.schema.fieldNames()]
         )
 
-        if self._is_hierarchical_bert:
-            return create_hierarchical_sequence_data(
-                person=self._dependency_dict[PERSON],
-                visit_occurrence=self._dependency_dict[VISIT_OCCURRENCE],
-                patient_events=cohort_ehr_records,
-                allow_measurement_only=self._allow_measurement_only,
-            )
-
         if self._is_feature_concept_frequency:
             return create_concept_frequency_data(cohort_ehr_records, None)
 
@@ -840,7 +858,6 @@ def create_prediction_cohort(
         aggregate_by_hour=spark_args.aggregate_by_hour,
         is_new_patient_representation=spark_args.is_new_patient_representation,
         gpt_patient_sequence=spark_args.gpt_patient_sequence,
-        is_hierarchical_bert=spark_args.is_hierarchical_bert,
         classic_bert_seq=spark_args.classic_bert_seq,
         is_first_time_outcome=spark_args.is_first_time_outcome,
         # If the outcome negative query exists, that means we need to remove those questionable
@@ -862,4 +879,5 @@ def create_prediction_cohort(
         meds_format=spark_args.meds_format,
         cache_events=spark_args.cache_events,
         should_construct_artificial_visits=spark_args.should_construct_artificial_visits,
+        duplicate_records=spark_args.duplicate_records,
     ).build()

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/const/common.py

@@ -7,6 +7,8 @@ DEVICE_EXPOSURE = "device_exposure"
 OBSERVATION = "observation"
 MEASUREMENT = "measurement"
 PROCESSED_MEASUREMENT = "processed_measurement"
+PROCESSED_OBSERVATION = "processed_observation"
+PROCESSED_DEVICE = "processed_device"
 CATEGORICAL_MEASUREMENT = "categorical_measurement"
 OBSERVATION_PERIOD = "observation_period"
 DEATH = "death"

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/decorators/clinical_event_decorator.py

@@ -1,8 +1,7 @@
-import os.path
-
 from ..const.common import (
     MEASUREMENT,
     CATEGORICAL_MEASUREMENT,
+    OBSERVATION,
 )
 from pyspark.sql import SparkSession, DataFrame, functions as F, Window as W, types as T
 
@@ -127,6 +126,13 @@ class ClinicalEventDecorator(PatientEventDecorator):
             ).otherwise(F.col("datetime"))
         )
 
+        outpatient_datetime_udf = (
+            F.when(
+                F.col("is_inpatient") == 0,
+                F.to_timestamp("date")
+            ).otherwise(F.col("datetime"))
+        )
+
         patient_events = (
             patient_events.join(visits, ["cohort_member_id", "visit_occurrence_id"])
             .withColumn("datetime", F.to_timestamp("datetime"))
@@ -136,6 +142,7 @@ class ClinicalEventDecorator(PatientEventDecorator):
             .withColumn("visit_end_datetime", F.expr("visit_end_datetime - INTERVAL 1 MINUTE"))
             .withColumn("date", bound_medical_event_date)
             .withColumn("datetime", bound_medical_event_datetime)
+            .withColumn("datetime", outpatient_datetime_udf)
             .withColumn("concept_order", concept_order_udf)
             .withColumn("visit_concept_order", visit_concept_order_udf)
             .drop("is_inpatient", "visit_end_date", "visit_end_datetime")
@@ -154,9 +161,9 @@ class ClinicalEventDecorator(PatientEventDecorator):
         # Create the concept_value_mask field to indicate whether domain values should be skipped
         # As of now only measurement has values, so other domains would be skipped.
         patient_events = patient_events.withColumn(
-            "concept_value_mask", (F.col("domain").isin(MEASUREMENT, CATEGORICAL_MEASUREMENT)).cast("int")
+            "concept_value_mask", (F.col("number_as_value").isNotNull() | F.col("concept_as_value").isNotNull()).cast("int")
         ).withColumn(
-            "is_numeric_type", (F.col("domain") == MEASUREMENT).cast("int")
+            "is_numeric_type", F.col("number_as_value").isNotNull().cast("int")
         ).withColumn(
             "mlm_skip_value",
             (F.col("domain").isin([MEASUREMENT, CATEGORICAL_MEASUREMENT])).cast("int"),

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/tools/extract_features.py

@@ -13,10 +13,12 @@ from cehrbert_data.decorators import AttType
 from cehrbert_data.utils.spark_parse_args import create_spark_args
 from cehrbert_data.utils.spark_utils import (
     extract_ehr_records,
+    preprocess_domain_table,
     construct_artificial_visits,
     create_sequence_data_with_att,
     create_concept_frequency_data,
 )
+from cehrbert_data.const.common import PERSON, VISIT_OCCURRENCE
 
 
 class PredictionType(Enum):
@@ -128,15 +130,41 @@ def main(args):
         ehr_records.write.mode("overwrite").parquet(ehr_records_temp_folder)
         ehr_records = spark.read.parquet(ehr_records_temp_folder)
 
-    visit_occurrence = spark.read.parquet(os.path.join(args.input_folder, "visit_occurrence"))
+
+    birthdate_udf = f.coalesce(
+        "birth_datetime",
+        f.concat("year_of_birth", f.lit("-01-01")).cast("timestamp"),
+    )
+
+    person = preprocess_domain_table(spark, args.input_folder, PERSON)
+    patient_demographic = person.select(
+        "person_id",
+        birthdate_udf.alias("birth_datetime"),
+        "race_concept_id",
+        "gender_concept_id",
+    )
+
+
+    visit_occurrence = preprocess_domain_table(spark, args.input_folder, VISIT_OCCURRENCE)
 
     if args.should_construct_artificial_visits:
         ehr_records, visit_occurrence = construct_artificial_visits(
             ehr_records,
             visit_occurrence=visit_occurrence,
             spark=spark,
-            persistence_folder = str(os.path.join(args.output_folder, args.cohort_name))
+            persistence_folder = str(os.path.join(args.output_folder, args.cohort_name)),
+            duplicate_records=args.duplicate_records
         )
+        # Update age if some of the ehr_records have been re-associated with the new visits
+        ehr_records = ehr_records.join(
+            patient_demographic.select("person_id", "birth_datetime"),
+            "person_id",
+        ).join(
+            visit_occurrence.select("visit_occurrence_id", "visit_start_date"), "visit_occurrence_id"
+        ).withColumn(
+            "age",
+            f.ceil(f.months_between(f.col("visit_start_date"), f.col("birth_datetime")) / f.lit(12))
+        ).drop("visit_start_date", "birth_datetime")
 
     cohort_visit_occurrence = visit_occurrence.join(
         cohort.select("person_id", "cohort_member_id", "index_date"),
@@ -157,22 +185,6 @@ def main(args):
             f.col("visit_end_date").cast(t.TimestampType()),
             f.col("visit_start_datetime")
         ).cast(t.TimestampType())
-    ).where(
-        f.col("visit_start_datetime") <=
-        f.expr(f"index_date - INTERVAL {args.hold_off_window} DAYS + INTERVAL 0.1 SECOND")
-    )
-
-
-    birthdate_udf = f.coalesce(
-        "birth_datetime",
-        f.concat("year_of_birth", f.lit("-01-01")).cast("timestamp"),
-    )
-    person = spark.read.parquet(os.path.join(args.input_folder, "person"))
-    patient_demographic = person.select(
-        "person_id",
-        birthdate_udf.alias("birth_datetime"),
-        "race_concept_id",
-        "gender_concept_id",
     )
 
     age_udf = f.ceil(f.months_between(f.col("visit_start_date"), f.col("birth_datetime")) / f.lit(12))

{cehrbert_data-0.0.7 → cehrbert_data-0.0.8}/src/cehrbert_data/utils/spark_parse_args.py

@@ -416,4 +416,10 @@ def create_spark_args(parse: bool = True):
         help="Indicate whether we should construct artificial visits for "
              "the problem list records that could occur years ahead",
     )
+    parser.add_argument(
+        "--duplicate_records",
+        dest="duplicate_records",
+        action="store_true",
+        help="Indicate whether we want to duplicate the problem list records when constructing artificial visits"
+    )
     return parser.parse_args() if parse else parser