cehrbert-data 0.1.0__tar.gz → 0.1.2__tar.gz

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cehrbert_data
-Version: 0.1.0
+Version: 0.1.2
 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
@@ -13,9 +13,9 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: numpy==1.24.3
-Requires-Dist: packaging==23.2
-Requires-Dist: pandas==2.2.0
+Requires-Dist: numpy>=1.24.3
+Requires-Dist: packaging>=23.2
+Requires-Dist: pandas>=2.2.0
 Requires-Dist: pyspark==3.1.2
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/pyproject.toml

@@ -26,9 +26,9 @@ classifiers = [
 ]
 
 dependencies = [
-    "numpy==1.24.3",
-    "packaging==23.2",
-    "pandas==2.2.0",
+    "numpy>=1.24.3",
+    "packaging>=23.2",
+    "pandas>=2.2.0",
     "pyspark==3.1.2"
 ]
 

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/apps/generate_included_concept_list.py

@@ -53,7 +53,11 @@ def main(
     The function processes patient event data across various domain tables, excludes low-frequency
     concepts, and saves the filtered concepts to a specified output folder.
     """
-    spark = SparkSession.builder.appName("Generate concept list").getOrCreate()
+    spark = (
+        SparkSession.builder.appName("Generate concept list")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
 
     # Exclude measurement from domain_table_list if exists because we need to process measurement
     # in a different way

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/apps/generate_training_data.py

@@ -54,7 +54,11 @@ def main(
         duplicate_records: bool = False,
         disconnect_problem_list_records: bool = False,
 ):
-    spark = SparkSession.builder.appName("Generate CEHR-BERT Training Data").getOrCreate()
+    spark = (
+        SparkSession.builder.appName("Generate CEHR-BERT Training Data")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
 
     logger = logging.getLogger(__name__)
     logger.info(
@@ -120,6 +124,7 @@ def main(
         "person_id",
         "discharged_to_concept_id",
     )
+
     person = preprocess_domain_table(spark, input_folder, PERSON)
     birth_datetime_udf = F.coalesce("birth_datetime", F.concat("year_of_birth", F.lit("-01-01")).cast("timestamp"))
     person = person.select(
@@ -128,22 +133,12 @@ def main(
         "race_concept_id",
         "gender_concept_id",
     )
-
     visit_occurrence_person = visit_occurrence.join(person, "person_id").withColumn(
         "age",
         F.ceil(F.months_between(F.col("visit_start_date"), F.col("birth_datetime")) / F.lit(12)),
     )
     visit_occurrence_person = visit_occurrence_person.drop("birth_datetime")
 
-    death = preprocess_domain_table(spark, input_folder, DEATH) if include_death else None
-
-    if include_concept_list and patient_ehr_events:
-        # Filter out concepts
-        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"
-        )
-
     patient_ehr_events = (
         patient_ehr_events.join(visit_occurrence_person, "visit_occurrence_id")
         .select(
@@ -153,10 +148,12 @@ def main(
         .withColumn("cohort_member_id", F.col("person_id"))
     )
 
-    # Apply the age security measure
-    # We only keep the patient records, whose corresponding age is less than 90
-    if apply_age_filter:
-        patient_ehr_events = patient_ehr_events.where(F.col("age") < 90)
+    if include_concept_list and patient_ehr_events:
+        # Filter out concepts
+        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 not continue_from_events:
         patient_ehr_events.write.mode("overwrite").parquet(os.path.join(output_folder, "all_patient_events"))
@@ -164,24 +161,21 @@ def main(
     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_ehr_events, visit_occurrence_person = construct_artificial_visits(
+        patient_ehr_events, visit_occurrence = construct_artificial_visits(
             patient_ehr_events,
-            visit_occurrence_person,
+            visit_occurrence,
             spark=spark,
             persistence_folder=output_folder,
             duplicate_records=duplicate_records,
             disconnect_problem_list_records=disconnect_problem_list_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")
+
+    death = preprocess_domain_table(spark, input_folder, DEATH) if include_death else None
+
+    # Apply the age security measure
+    # We only keep the patient records, whose corresponding age is less than 90
+    if apply_age_filter:
+        patient_ehr_events = patient_ehr_events.where(F.col("age") < 90)
 
     if is_new_patient_representation:
         patient_sequence_data = create_sequence_data_with_att(

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/cohorts/spark_app_base.py

@@ -137,7 +137,11 @@ class BaseCohortBuilder(ABC):
         # Validate if the data folders exist
         validate_date_folder(self._input_folder, self._query_builder.get_dependency_list())
 
-        self.spark = SparkSession.builder.appName(f"Generate {self._query_builder.get_cohort_name()}").getOrCreate()
+        self.spark = (
+            SparkSession.builder.appName(f"Generate {self._query_builder.get_cohort_name()}")
+            .config("spark.sql.session.timeZone", "UTC")
+            .getOrCreate()
+        )
 
         self._dependency_dict = instantiate_dependencies(
             self.spark, self._input_folder, self._query_builder.get_dependency_list()
@@ -409,7 +413,11 @@ class NestedCohortBuilder:
             f"disconnect_problem_list_records: {disconnect_problem_list_records}\n"
         )
 
-        self.spark = SparkSession.builder.appName(f"Generate {self._cohort_name}").getOrCreate()
+        self.spark = (
+            SparkSession.builder.appName(f"Generate {self._cohort_name}")
+            .config("spark.sql.session.timeZone", "UTC")
+            .getOrCreate()
+        )
         self._dependency_dict = instantiate_dependencies(self.spark, self._input_folder, DEFAULT_DEPENDENCY)
 
         # Validate the input and output folders

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/connect_omop_visit.py

@@ -226,7 +226,11 @@ def step_2_connect_outpatient_to_inpatient(
 
 
 def main(args):
-    spark = SparkSession.builder.appName("Clean up visit_occurrence").getOrCreate()
+    spark = (
+        SparkSession.builder.appName("Clean up visit_occurrence")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
     visit_occurrence = spark.read.parquet(os.path.join(args.input_folder, "visit_occurrence"))
     visit_occurrence_step_1, in_to_in_visit_mapping = step_1_consolidate_inpatient_visits(
         spark,

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/convert_prediction_time_to_str.py

@@ -26,9 +26,10 @@ def convert_file(file_pair):
         # Read with Polars
         df = pl.read_parquet(input_file)
 
-        # Convert datetime to string (ISO format)
+        # Convert datetime to string (ISO format) in UTC to avoid timezone issues
         df = df.with_columns([
             pl.col("prediction_time")
+            .dt.convert_time_zone("UTC")
             .cast(pl.Datetime("us"))
             .dt.strftime("%Y-%m-%d %H:%M:%S.%f")
             .alias("prediction_time")

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/download_omop_tables.py

@@ -107,7 +107,11 @@ if __name__ == "__main__":
     )
 
     ARGS = parser.parse_args()
-    spark = SparkSession.builder.appName("Download OMOP tables").getOrCreate()
+    spark = (
+        SparkSession.builder.appName("Download OMOP tables")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
     domain_table_list = ARGS.domain_table_list
     credential_path = ARGS.credential_path
     download_folder = ARGS.output_folder

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/ehrshot_to_omop.py

@@ -784,7 +784,11 @@ def drop_duplicate_visits(data: DataFrame) -> DataFrame:
 
 
 def main(args):
-    spark = SparkSession.builder.appName("Convert EHRShot Data").getOrCreate()
+    spark = (
+        SparkSession.builder.appName("Convert EHRShot Data")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
 
     logger.info(
         f"ehr_shot_file: {args.ehr_shot_file}\n"

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/extract_features.py

@@ -71,6 +71,7 @@ def main(args):
         "spark.sql.legacy.parquet.int96RebaseModeInWrite",
         "CORRECTED"
     )
+    spark.conf.set("spark.sql.session.timeZone", "UTC")
 
     cohort_dir = os.path.expanduser(args.cohort_dir)
     is_parquet = False

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/sample_omop_tables.py

@@ -23,6 +23,7 @@ def main(args):
         .config("spark.sql.legacy.parquet.int96RebaseModeInWrite", "CORRECTED")
         .config("spark.sql.legacy.parquet.datetimeRebaseModeInRead", "CORRECTED")
         .config("spark.sql.legacy.parquet.datetimeRebaseModeInWrite", "CORRECTED")
+        .config("spark.sql.session.timeZone", "UTC")
         .getOrCreate()
     )
     patient_sample = spark.read.parquet(args.person_sample)

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/tools/update_omop_visit.py

@@ -7,7 +7,11 @@ from pyspark.sql import SparkSession
 from pyspark.sql import functions as f
 
 def main(args):
-    spark = SparkSession.builder.appName("Clean up visit_occurrence").getOrCreate()
+    spark = (
+        SparkSession.builder.appName("Clean up visit_occurrence")
+        .config("spark.sql.session.timeZone", "UTC")
+        .getOrCreate()
+    )
     visit_mapping = spark.read.parquet(
         os.path.join(args.output_folder, "visit_mapping")
     )

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data/utils/spark_utils.py

@@ -44,7 +44,7 @@ DOMAIN_KEY_FIELDS = {
             "condition_concept_id",
             "condition_start_date",
             "condition_start_datetime",
-            "condition"
+            "condition_occurrence"
         )
     ],
     "procedure_occurrence_id": [
@@ -52,7 +52,7 @@ DOMAIN_KEY_FIELDS = {
             "procedure_concept_id",
             "procedure_date",
             "procedure_datetime",
-            "procedure"
+            "procedure_occurrence"
         )
     ],
     "drug_exposure_id": [
@@ -60,7 +60,7 @@ DOMAIN_KEY_FIELDS = {
             "drug_concept_id",
             "drug_exposure_start_date",
             "drug_exposure_start_datetime",
-            "drug"
+            "drug_exposure"
         )
     ],
     "measurement_id": [
@@ -84,13 +84,13 @@ DOMAIN_KEY_FIELDS = {
             "device_concept_id",
             "device_exposure_start_date",
             "device_exposure_start_datetime",
-            "device"
+            "device_exposure"
         )
     ],
     "death_date": [("cause_concept_id", "death_date", "death_datetime", "death")],
     "visit_concept_id": [
-        ("visit_concept_id", "visit_start_date", "visit"),
-        ("discharged_to_concept_id", "visit_end_date", "visit"),
+        ("visit_concept_id", "visit_start_date", "visit_occurrence"),
+        ("discharged_to_concept_id", "visit_end_date", "visit_occurrence"),
     ],
 }
 
@@ -180,7 +180,7 @@ def extract_events_by_domain(
     ) in get_key_fields(domain_table):
 
         if is_domain_numeric(domain_table_name):
-            concept = kwargs.get("concept")
+            concept: DataFrame = kwargs.get("concept")
             spark = kwargs.get("spark", None)
             persistence_folder = kwargs.get("persistence_folder", None)
             refresh = kwargs.get("refresh_measurement", False)
@@ -210,7 +210,7 @@ def extract_events_by_domain(
             domain_records = domain_table.where(F.col(date_field).isNotNull()).where(
                 F.col(concept_id_field).isNotNull()
             )
-            datetime_field_udf = F.to_timestamp(F.coalesce(datetime_field, date_field), "yyyy-MM-dd HH:mm:ss")
+            datetime_field_udf = F.to_timestamp(F.coalesce(datetime_field, date_field))
             domain_records = (
                 domain_records.where(F.col(concept_id_field).cast("string") != "0")
                 .withColumn("date", F.to_date(F.col(date_field)))
@@ -222,17 +222,13 @@ def extract_events_by_domain(
                 domain_records["date"].cast("date"),
                 domain_records["datetime"].cast(T.TimestampType()),
                 domain_records["visit_occurrence_id"],
-                F.lit(domain_table_name).alias("domain"),
+                F.lit(domain_table_name.split("_")[0]).alias("domain"),
                 F.lit(None).cast("string").alias("event_group_id"),
                 F.lit(None).cast("float").alias("number_as_value"),
                 F.lit(None).cast("string").alias("concept_as_value"),
                 F.col("unit") if domain_has_unit(domain_records) else F.lit(NA).alias("unit"),
             ).distinct()
 
-            # Remove "Patient Died" from condition_occurrence
-            if domain_table_name == "condition_occurrence":
-                domain_records = domain_records.where("condition_concept_id != 4216643")
-
         if ehr_events is None:
             ehr_events = domain_records
         else:

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cehrbert_data
-Version: 0.1.0
+Version: 0.1.2
 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
@@ -13,9 +13,9 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: numpy==1.24.3
-Requires-Dist: packaging==23.2
-Requires-Dist: pandas==2.2.0
+Requires-Dist: numpy>=1.24.3
+Requires-Dist: packaging>=23.2
+Requires-Dist: pandas>=2.2.0
 Requires-Dist: pyspark==3.1.2
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"

{cehrbert_data-0.1.0 → cehrbert_data-0.1.2}/src/cehrbert_data.egg-info/requires.txt

@@ -1,6 +1,6 @@
-numpy==1.24.3
-packaging==23.2
-pandas==2.2.0
+numpy>=1.24.3
+packaging>=23.2
+pandas>=2.2.0
 pyspark==3.1.2
 
 [dev]