move demographic/birth factors into inputs data pipeline

databricks_workflows/nhp_data-extract_nhp_for_containers.yaml

@@ -81,43 +81,9 @@ resources:
               job_cluster_key: nhp_data
               libraries:
                 - whl: ../dist/*.whl
-        - task_key: run_extract_demographic_factors_data
-          depends_on:
-            - task_key: run_extract_ecds_data
-          for_each_task:
-            inputs: "{{job.parameters.years}}"
-            task:
-              task_key: extract_demographic_factors_data
-              python_wheel_task:
-                package_name: nhp_data
-                entry_point: model_data-demographic_factors
-                parameters:
-                  - "{{job.parameters.data_version}}"
-                  - "{{input}}"
-                  - "{{job.parameters.projection_year}}"
-              job_cluster_key: nhp_data
-              libraries:
-                - whl: ../dist/*.whl
-        - task_key: run_extract_birth_factors_data
-          depends_on:
-            - task_key: run_extract_demographic_factors_data
-          for_each_task:
-            inputs: "{{job.parameters.years}}"
-            task:
-              task_key: extract_birth_factors_data
-              python_wheel_task:
-                package_name: nhp_data
-                entry_point: model_data-birth_factors
-                parameters:
-                  - "{{job.parameters.data_version}}"
-                  - "{{input}}"
-                  - "{{job.parameters.projection_year}}"
-              job_cluster_key: nhp_data
-              libraries:
-                - whl: ../dist/*.whl
         - task_key: run_extract_inequalities_data
           depends_on:
-            - task_key: run_extract_birth_factors_data
+            - task_key: run_extract_ecds_data
           for_each_task:
             inputs: "{{job.parameters.years}}"
             task:

databricks_workflows/nhp_data-inputs-provider.yaml

@@ -86,6 +86,17 @@ resources:
           job_cluster_key: nhp_data
           libraries:
             - whl: ../dist/*.whl
+        - task_key: extract_demographics
+          depends_on:
+            - task_key: extract_inequalities
+          python_wheel_task:
+            package_name: nhp_data
+            entry_point: inputs_data-demographics
+            parameters:
+              - provider
+          job_cluster_key: nhp_data
+          libraries:
+            - whl: ../dist/*.whl
       job_clusters:
         - job_cluster_key: nhp_data
           new_cluster:

databricks_workflows/nhp_data-reference_data.yaml

@@ -83,6 +83,17 @@ resources:
             parameters:
               - "{{job.parameters.day_procedures_base_year}}"
           job_cluster_key: nhp_data
+          libraries:
+            - whl: ../dist/*.whl
+        - task_key: create_population_provider
+          depends_on:
+            - task_key: create_day_procedure_code_list
+          python_wheel_task:
+            package_name: nhp_data
+            entry_point: reference-population_provider
+          job_cluster_key: nhp_data
+          libraries:
+            - whl: ../dist/*.whl
       job_clusters:
         - job_cluster_key: nhp_data
           new_cluster:

pyproject.toml

@@ -74,6 +74,7 @@ inputs_data-expat_repat = "nhp.data.inputs_data.expat_repat:main"
 inputs_data-procedures = "nhp.data.inputs_data.procedures:main"
 inputs_data-rates = "nhp.data.inputs_data.rates:main"
 inputs_data-inequalities = "nhp.data.inputs_data.inequalities:main"
+inputs_data-demographics = "nhp.data.inputs_data.demographics:main"
 
 raw_data-aae = "nhp.data.raw_data.aae:main"
 raw_data-ecds = "nhp.data.raw_data.ecds:main"
@@ -90,8 +91,6 @@ model_data-health_status_adjustment-generate_provider_gams = "nhp.data.model_dat
 model_data-health_status_adjustment-generate_icb_gams = "nhp.data.model_data.health_status_adjustment.generate_icb_gams:main"
 model_data-health_status_adjustment-generate_national_gams = "nhp.data.model_data.health_status_adjustment.generate_national_gams:main"
 model_data-aae = "nhp.data.model_data.aae:main"
-model_data-birth_factors = "nhp.data.model_data.birth_factors:main"
-model_data-demographic_factors = "nhp.data.model_data.demographic_factors:main"
 model_data-inequalities = "nhp.data.model_data.inequalities:main"
 model_data-ip = "nhp.data.model_data.ip:main"
 model_data-ip-tpmas = "nhp.data.model_data.ip_tpmas:main"
@@ -108,5 +107,6 @@ reference-lsoa_lookups = "nhp.data.reference.lsoa_lookups:main"
 reference-population_by_lsoa21 = "nhp.data.reference.population_by_lsoa21:main"
 reference-provider_catchments = "nhp.data.reference.provider_catchments:main"
 reference-icb_catchments = "nhp.data.reference.icb_catchments:main"
+reference-population_provider = "nhp.data.reference.population_provider:main"
 
 extract_data = "nhp.data.extract_data:main"

src/nhp/data/aggregated_data/ecds.py

@@ -3,7 +3,7 @@
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql import functions as F
 
-from nhp.data.get_spark import get_spark
+from nhp.data.get_spark import PartitionOverwriteMode, get_spark
 from nhp.data.table_names import table_names
 
 
@@ -42,7 +42,7 @@ def get_ecds_data(spark: SparkSession) -> DataFrame:
     )
 
 
-def generate_ecds_data(spark: SparkSession, ecds: DataFrame) -> None:
+def generate_ecds_data(ecds: DataFrame) -> None:
     """Generate ECDS Data"""
     (
         ecds.withColumn("index", F.expr("uuid()"))
@@ -55,6 +55,6 @@ def generate_ecds_data(spark: SparkSession, ecds: DataFrame) -> None:
 
 def main() -> None:
     """main method"""
-    spark = get_spark()
+    spark = get_spark(PartitionOverwriteMode.STATIC)
     ecds = get_ecds_data(spark)
-    generate_ecds_data(spark, ecds)
+    generate_ecds_data(ecds)

src/nhp/data/aggregated_data/outpatients.py

@@ -3,7 +3,7 @@
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql import functions as F
 
-from nhp.data.get_spark import get_spark
+from nhp.data.get_spark import PartitionOverwriteMode, get_spark
 from nhp.data.table_names import table_names
 
 
@@ -45,7 +45,7 @@ def get_outpatients_data(spark: SparkSession) -> DataFrame:
     )
 
 
-def generate_outpatients_data(spark: SparkSession, opa: DataFrame) -> None:
+def generate_outpatients_data(opa: DataFrame) -> None:
     """Generate Outpatients Data"""
     (
         opa.write.partitionBy("fyear", "provider")
@@ -57,6 +57,6 @@ def generate_outpatients_data(spark: SparkSession, opa: DataFrame) -> None:
 
 def main() -> None:
     """main method"""
-    spark = get_spark()
+    spark = get_spark(PartitionOverwriteMode.STATIC)
     opa = get_outpatients_data(spark)
-    generate_outpatients_data(spark, opa)
+    generate_outpatients_data(opa)

src/nhp/data/get_spark.py

@@ -1,9 +1,18 @@
 """Helper method to get the spark session."""
 
+import enum
+
 from pyspark.sql import SparkSession
 
 
-def get_spark() -> SparkSession:
+class PartitionOverwriteMode(enum.Enum):
+    DYNAMIC = "dynamic"
+    STATIC = "static"
+
+
+def get_spark(
+    partition_overwrite_mode: PartitionOverwriteMode = PartitionOverwriteMode.DYNAMIC,
+) -> SparkSession:
     """Get spark session
 
     :return: get the spark session to use
@@ -14,5 +23,7 @@ def get_spark() -> SparkSession:
     from databricks.connect import DatabricksSession
 
     spark = DatabricksSession.builder.getOrCreate()
-    spark.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")
+    spark.conf.set(
+        "spark.sql.sources.partitionOverwriteMode", partition_overwrite_mode.value
+    )
     return spark

src/nhp/data/inputs_data/demographics.py

@@ -0,0 +1,62 @@
+"""Demographics/Birth/Catchments"""
+
+import sys
+
+import pyspark.sql.functions as F
+from pyspark.sql import SparkSession
+
+from nhp.data.get_spark import get_spark
+from nhp.data.inputs_data.save_parquet import save_parquet
+from nhp.data.table_names import table_names
+
+
+def extract_demographic_factors(path: str, spark: SparkSession) -> None:
+    """Extract Demographic Factors data
+
+    :param path: the path to save the extracted data
+    :type path: str
+    :param spark: the spark session to use
+    :type spark: SparkSession
+    """
+    df = spark.read.table(table_names.reference_population_provider_demographics)
+    save_parquet(df, f"{path}/demographic_factors")
+
+
+def extract_birth_factors(path: str, spark: SparkSession) -> None:
+    """Extract Birth Factors data
+
+    :param path: the path to save the extracted data
+    :type path: str
+    :param spark: the spark session to use
+    :type spark: SparkSession
+    """
+    df = spark.read.table(table_names.reference_population_provider_births)
+    save_parquet(df, f"{path}/birth_factors")
+
+
+def extract_provider_catchments(path: str, spark: SparkSession) -> None:
+    """Extract Provider Catchments data
+
+    :param path: the path to save the extracted data
+    :type path: str
+    :param spark: the spark session to use
+    :type spark: SparkSession
+    """
+    df = (
+        spark.read.table(table_names.reference_provider_lad23_splits)
+        .filter(F.col("fyear") >= 202324)
+        .orderBy("fyear", "provider", "lad23cd", "sex", "age")
+    )
+    save_parquet(df, f"{path}/provider_catchments")
+
+
+def main() -> None:
+    geography_column = sys.argv[1]
+    assert geography_column in ["provider", "lad23cd"], "invalid geography_column"
+
+    path = f"{table_names.inputs_save_path}/{geography_column}"
+
+    spark = get_spark()
+    extract_demographic_factors(path, spark)
+    extract_birth_factors(path, spark)
+    extract_provider_catchments(path, spark)

src/nhp/data/model_data/generate_synthetic_data.py

@@ -75,8 +75,6 @@ def generate(self) -> None:
         self._inequalities()
         self._aae()
         self._op()
-        self._birth_factors()
-        self._demographic_factors()
         self._hsa_activity_tables()
 
     @property
@@ -392,26 +390,6 @@ def _op(self, df: DataFrame) -> pd.DataFrame:
 
         return op
 
-    @generate_data("birth_factors")
-    def _birth_factors(self, df: DataFrame) -> pd.DataFrame:
-        return (
-            df.drop("dataset")
-            .filter(~F.col("variant").startswith("custom_projection_"))
-            .toPandas()
-            .groupby(["variant", "sex", "age"], as_index=False)
-            .mean()
-        )
-
-    @generate_data("demographic_factors")
-    def _demographic_factors(self, df: DataFrame) -> pd.DataFrame:
-        return (
-            df.drop("dataset")
-            .filter(~F.col("variant").startswith("custom_projection_"))
-            .toPandas()
-            .groupby(["variant", "sex", "age"], as_index=False)
-            .mean()
-        )
-
     @generate_data("hsa_activity_tables")
     def _hsa_activity_tables(self, df: DataFrame) -> pd.DataFrame:
         return (