sws-spark-dissemination-helper 0.0.93__py3-none-any.whl → 0.0.183__py3-none-any.whl

sws_spark_dissemination_helper/SWSBronzeIcebergSparkHelper.py

@@ -1,6 +1,7 @@
 import logging
+import time
 from copy import copy
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
 
 import pyspark.sql.functions as F
 from pyspark.sql import DataFrame, SparkSession
@@ -25,6 +26,9 @@ class SWSBronzeIcebergSparkHelper:
         domain_code: str,
         dataset_details: dict = None,
         dataset_tables: DatasetTables = None,
+        keep_history: bool = False,
+        write_csv: bool = True,
+        source_tag: Union[str, None] = None,
     ) -> None:
         self.spark: SparkSession = spark
         self.dataset_details: dict = dataset_details
@@ -35,6 +39,9 @@ class SWSBronzeIcebergSparkHelper:
         self.dataset_tables: DatasetTables = dataset_tables
         self.iceberg_tables: IcebergTables = iceberg_tables
         self.domain_code = domain_code
+        self.keep_history: bool = keep_history
+        self.write_csv: bool = write_csv
+        self.source_tag: Union[str, None] = source_tag
 
         if dataset_details is not None:
             (
@@ -82,6 +89,7 @@ class SWSBronzeIcebergSparkHelper:
                 self.df_obs_coord,
                 self.df_metadata,
                 self.df_meta_elem,
+                self.df_tag_observation,
             ) = self.raw_data
 
             (
@@ -91,10 +99,11 @@ class SWSBronzeIcebergSparkHelper:
                 self.df_meta_elem_type,
                 self.df_language,
                 self.df_unit_of_measure,
+                self.df_dataset,
                 self.dfs_dimension,
             ) = self.raw_reference_data
 
-            self.df_user = self.raw_operational_data
+            (self.df_user, self.df_tag) = self.raw_operational_data
 
     def _get_dim_time_flag_columns(self) -> Tuple[List[str], List[str], str, List[str]]:
         """Extract the dimension columns with time, without time, the time column and the flag columns names."""
@@ -149,7 +158,7 @@ class SWSBronzeIcebergSparkHelper:
 
         return dfs_dimension
 
-    def _prepare_element_uom(self) -> DataFrame:
+    def _prepare_element_uom(self) -> Union[DataFrame, None]:
         """Prepare the element and unit of measure join."""
 
         # Get the element DataFrame
@@ -161,23 +170,24 @@ class SWSBronzeIcebergSparkHelper:
             if dimension_column == self.element_column
         )
 
-        # Join the element and the unit_of_measure
-        df_element_uom = (
-            df_element.alias("e")
-            .join(
-                self.df_unit_of_measure.alias("u"),
-                col("e.unit_of_measure") == col("u.id"),
-            )
-            .select(
-                col("e.code").alias("element_code"),
-                col("u.code").alias("unit_of_measure"),
-                col("u.symbol").alias("unit_of_measure_symbol"),
-                col("u.base_unit").alias("unit_of_measure_base_unit"),
-                col("u.multiplier").alias("unit_of_measure_multiplier"),
+        if any("unit_of_measure" == column.lower() for column in df_element.columns):
+            # Join the element and the unit_of_measure
+            df_element_uom = (
+                df_element.alias("e")
+                .join(
+                    self.df_unit_of_measure.alias("u"),
+                    col("e.unit_of_measure") == col("u.id"),
+                )
+                .select(
+                    col("e.code").alias("element_code"),
+                    col("u.code").alias("unit_of_measure"),
+                    col("u.symbol").alias("unit_of_measure_symbol"),
+                    col("u.base_unit").alias("unit_of_measure_base_unit"),
+                    col("u.multiplier").alias("unit_of_measure_multiplier"),
+                )
             )
-        )
 
-        return df_element_uom
+            return df_element_uom
 
     def _gen_denormalized_observation(self) -> DataFrame:
         """Original query upon which the below computation is based
@@ -269,20 +279,170 @@ class SWSBronzeIcebergSparkHelper:
                 .withColumnRenamed("code", dimension_column)
             )
 
-        df_intermediate = (
-            df_intermediate.alias("d")
-            .join(
-                F.broadcast(df_element_uom).alias("e"),
-                col(f"d.{self.element_column}") == col("e.element_code"),
-                "left",
+        if df_element_uom is not None:
+            df_intermediate = (
+                df_intermediate.alias("d")
+                .join(
+                    F.broadcast(df_element_uom).alias("e"),
+                    col(f"d.{self.element_column}") == col("e.element_code"),
+                    "left",
+                )
+                .drop("element_code")
             )
-            .drop("element_code")
-        )
 
         df_obs_denorm = df_intermediate
 
         return df_obs_denorm
 
+    def _gen_denormalized_observation_sql(self) -> DataFrame:
+        # ----------------
+        # Prepare dataframes for the joins
+        # ----------------
+
+        select_statement = """
+            o.id,
+            o.value,
+            u.email,
+            o.created_on,
+            o.replaced_on,
+            o.version"""
+
+        from_statement = f"""
+        FROM {self.dataset_tables.OBSERVATION.iceberg_id} o
+            JOIN {self.dataset_tables.USER.iceberg_id} u ON u.id = o.created_by
+            LEFT JOIN {self.dataset_tables.OBSERVATION_COORDINATE.iceberg_id} AS oc ON oc.id = o.observation_coordinates"""
+
+        hint_statement = ""
+
+        id_to_flag_col_mapping = {v: k for k, v in self.flag_col_to_id_mapping.items()}
+        for flag_col in self.flag_columns:
+            select_statement += f",\no.{id_to_flag_col_mapping[flag_col]} AS {flag_col}"
+
+        id_to_dim_col_mapping = {v: k for k, v in self.dim_col_to_id_mapping.items()}
+        for i, (dim_col, cl) in enumerate(
+            zip(self.dim_columns_w_time, self.dataset_tables.CODELISTS)
+        ):
+            select_statement += f",\nd{i}.code AS {dim_col}"
+            from_statement += f"\nLEFT JOIN {cl.iceberg_id} d{i} ON d{i}.id = oc.{id_to_dim_col_mapping[dim_col]}"
+            hint_statement = (
+                hint_statement + f", BROADCAST({cl.iceberg_id})"
+                if hint_statement
+                else f"BROADCAST({cl.iceberg_id})"
+            )
+
+        hint_statement = "/*+ " + hint_statement + " */"
+
+        final_query = "SELECT " + hint_statement + select_statement + from_statement
+        if not self.keep_history:
+            final_query += "\nWHERE o.replaced_on IS NULL"
+
+        logging.info("Final query for merging observation and observation_coordinares")
+        logging.info(final_query)
+
+        df_obs_denorm = self.spark.sql(final_query)
+
+        df_element_uom = self._prepare_element_uom()
+
+        dfs_dimension_w_validity = self._convert_dim_start_end_date_to_data()
+
+        # Join all the dimension codelists
+        for dimension_column, df_dimension in zip(
+            self.dim_columns_w_time, dfs_dimension_w_validity
+        ):
+            logging.debug(f"Joining dimension column: {dimension_column}")
+            logging.debug(f"df_obs_denorm columns: {df_obs_denorm.columns}")
+            logging.debug(
+                f"Is dimension {dimension_column} in the dataframe? {dimension_column in df_obs_denorm.columns}"
+            )
+            df_obs_denorm = (
+                df_obs_denorm.alias("o")
+                .join(
+                    F.broadcast(df_dimension.withColumnRenamed("id", "join_id")).alias(
+                        "d"
+                    ),
+                    col(f"{dimension_column}") == col("d.code"),
+                )
+                .drop("code", "join_id")
+            )
+            logging.debug(f"After join count: {df_obs_denorm.count()}")
+
+        if df_element_uom is not None:
+            df_obs_denorm = (
+                df_obs_denorm.alias("d")
+                .join(
+                    F.broadcast(df_element_uom).alias("e"),
+                    col(f"d.{self.element_column}") == col("e.element_code"),
+                    "left",
+                )
+                .drop("element_code")
+            )
+            logging.debug(f"After uom count: {df_obs_denorm.count()}")
+
+        return df_obs_denorm
+
+    def _gen_denormalized_observation_sql_from_tag(self) -> DataFrame:
+        # ----------------
+        # Prepare dataframes for the joins
+        # ----------------
+
+        select_statement = """
+            o.id,
+            o.value,
+            u.email,
+            o.created_on,
+            o.replaced_on,
+            o.version"""
+
+        from_statement = f"""
+        FROM {self.dataset_tables.OBSERVATION.iceberg_id} o
+            INNER JOIN {self.dataset_tables.TAG_OBSERVATION.iceberg_id} to ON o.id = to.observation
+            INNER JOIN {self.dataset_tables.TAG.iceberg_id} t ON to.tag = t.id
+            INNER JOIN {self.dataset_tables.DATASET.iceberg_id} d ON t.dataset = d.id
+            LEFT JOIN {self.dataset_tables.USER.iceberg_id} u ON u.id = o.created_by
+            LEFT JOIN {self.dataset_tables.OBSERVATION_COORDINATE.iceberg_id} AS oc ON oc.id = o.observation_coordinates"""
+
+        hint_statement = ""
+
+        id_to_flag_col_mapping = {v: k for k, v in self.flag_col_to_id_mapping.items()}
+        for flag_col in self.flag_columns:
+            select_statement += f",\no.{id_to_flag_col_mapping[flag_col]} AS {flag_col}"
+
+        id_to_dim_col_mapping = {v: k for k, v in self.dim_col_to_id_mapping.items()}
+        for i, (dim_col, cl) in enumerate(
+            zip(self.dim_columns_w_time, self.dataset_tables.CODELISTS)
+        ):
+            select_statement += f",\nd{i}.code AS {dim_col}"
+            from_statement += f"\nLEFT JOIN {cl.iceberg_id} d{i} ON d{i}.id = oc.{id_to_dim_col_mapping[dim_col]}"
+            hint_statement = (
+                hint_statement + f", BROADCAST({cl.iceberg_id})"
+                if hint_statement
+                else f"BROADCAST({cl.iceberg_id})"
+            )
+
+        hint_statement = "/*+ " + hint_statement + " */"
+
+        # TODO Add tag name as a parameter
+        where_statement = (
+            f"\nWHERE t.name = '{self.source_tag}' AND d.xml_name = '{self.dataset_id}'"
+        )
+
+        final_query = (
+            "SELECT "
+            + hint_statement
+            + select_statement
+            + from_statement
+            + where_statement
+        )
+        if not self.keep_history:
+            final_query += "\n AND o.replaced_on IS NULL"
+
+        logging.info("Final query for merging observation and observation_coordinares")
+        logging.info(final_query)
+
+        df_obs_denorm = self.spark.sql(final_query)
+
+        return df_obs_denorm
+
     def _gen_denormalized_metadata(self) -> DataFrame:
         """Original query upon which the below computation is based
 
@@ -346,6 +506,32 @@ class SWSBronzeIcebergSparkHelper:
 
         return df_meta_denorm
 
+    def _gen_denormalized_metadata_sql(self) -> DataFrame:
+        # ----------------
+        # Generate denormalized observation table
+        # ----------------
+
+        logging.info("meta_denorm start")
+
+        df_meta_denorm = self.spark.sql(
+            f"""
+        select m.observation as observation_id,
+            mt.code as type,
+            met.code as element_type,
+            l.country_code as language,
+            me.value
+        from {self.dataset_tables.METADATA_ELEMENT.iceberg_id} me
+            left join {self.dataset_tables.METADATA.iceberg_id} m on m.id = me.metadata
+            left join {self.dataset_tables.METADATA_ELEMENT_TYPE.iceberg_id} met on met.id = me.metadata_element_type
+            left join {self.dataset_tables.METADATA_TYPE.iceberg_id} mt on mt.id = m.metadata_type
+            left join {self.dataset_tables.LANGUAGE.iceberg_id} l on l.id = m.language
+        """
+        )
+
+        logging.info("meta_denorm write")
+
+        return df_meta_denorm
+
     def _gen_grouped_metadata(self) -> DataFrame:
         return (
             self._gen_denormalized_metadata()
@@ -366,6 +552,26 @@ class SWSBronzeIcebergSparkHelper:
             .agg(F.collect_list("metadata").alias("metadata"))
         )
 
+    def _gen_grouped_metadata_sql(self) -> DataFrame:
+        return (
+            self._gen_denormalized_metadata_sql()
+            .select(
+                col("observation_id"),
+                F.create_map(
+                    lit("type"),
+                    col("type"),
+                    lit("element_type"),
+                    col("element_type"),
+                    lit("language"),
+                    col("language"),
+                    lit("value"),
+                    col("value"),
+                ).alias("metadata"),
+            )
+            .groupby("observation_id")
+            .agg(F.collect_list("metadata").alias("metadata"))
+        )
+
     def _gen_bronze_data(self) -> DataFrame:
         return (
             self._gen_denormalized_observation()
@@ -378,9 +584,37 @@ class SWSBronzeIcebergSparkHelper:
             .drop("m.observation_id")
         )
 
+    def _gen_bronze_data_sql(self) -> DataFrame:
+        return (
+            self._gen_denormalized_observation_sql()
+            .alias("o")
+            .join(
+                self._gen_grouped_metadata_sql().alias("m"),
+                col("o.id") == col("m.observation_id"),
+                "left",
+            )
+            .drop("m.observation_id")
+        )
+
+    def _gen_bronze_data_sql_from_tag(self) -> DataFrame:
+        return (
+            self._gen_denormalized_observation_sql_from_tag()
+            .alias("o")
+            .join(
+                self._gen_grouped_metadata_sql().alias("m"),
+                col("o.id") == col("m.observation_id"),
+                "left",
+            )
+            .drop("m.observation_id")
+        )
+
     # TODO decouple data generation and data writing
-    def write_bronze_data_to_iceberg_and_csv(self) -> DataFrame:
-        self.df_bronze = self._gen_bronze_data()
+    def write_bronze_data_to_iceberg_and_csv(self, sql=True) -> DataFrame:
+
+        if sql:
+            self.df_bronze = self._gen_bronze_data_sql()
+        else:
+            self.df_bronze = self._gen_bronze_data()
 
         self.df_bronze.writeTo(self.iceberg_tables.BRONZE.iceberg_id).createOrReplace()
 
@@ -416,11 +650,13 @@ class SWSBronzeIcebergSparkHelper:
             description="Bronze table containing all the raw data imported from the SWS and denormalized",
             layer=TableLayer.BRONZE,
             private=True,
+            debug=True,
             type=TableType.ICEBERG,
             database=IcebergDatabases.BRONZE_DATABASE,
             table=self.iceberg_tables.BRONZE.table,
             path=self.iceberg_tables.BRONZE.path,
             structure={"columns": self.df_bronze.schema.jsonValue()["fields"]},
+            pinned_columns=[*self.dim_columns_w_time, "value", *self.flag_columns],
         )
         tag = tags.add_dissemination_table(
             self.dataset_id, self.tag_name, new_iceberg_table
@@ -433,6 +669,7 @@ class SWSBronzeIcebergSparkHelper:
             description="Bronze table containing all the raw data imported from the SWS and denormalized cached in csv",
             layer=TableLayer.BRONZE,
             private=True,
+            debug=True,
             type=TableType.CSV,
             path=self.iceberg_tables.BRONZE.csv_path,
             structure={"columns": self.df_bronze.schema.jsonValue()["fields"]},
@@ -445,34 +682,34 @@ class SWSBronzeIcebergSparkHelper:
         logging.info("Bronze Dissemination tags successfully written")
 
     def write_bronze_disseminated_tag_data_to_iceberg_and_csv(
-        self, dimensions: Dict[str, List[str]]
+        self, dimensions: Dict[str, List[str]] = {}, from_tag=False
     ) -> DataFrame:
 
-        create_branch_query = f"ALTER TABLE {self.iceberg_tables.BRONZE.iceberg_id}.`tag_{self.tag_name}` CREATE OR REPLACE BRANCH `diss_tag_{self.tag_name}`"  # AS OF VERSION `{tag_name}`
-        logging.info(f"create_branch_query: {create_branch_query}")
-        create_branch_query_result = self.spark.sql(create_branch_query).collect()
-
-        logging.info(f"result of create_branch_query: {create_branch_query_result}")
-
-        logging.info(f"dimensions: {dimensions}")
-        for dimension_name, codes in dimensions.items():
-            logging.info(f"dimension_name: {dimension_name}")
-            logging.info(f"codes: {codes}")
-            if len(codes) != 0:
-                not_in_codes = ",".join([f"'{code}'" for code in codes])
-                delete_from_branch_query = f"DELETE FROM {self.iceberg_tables.BRONZE.iceberg_id}.`branch_diss_tag_{self.tag_name}` WHERE {dimension_name} NOT IN ({not_in_codes})"
-                logging.info(f"delete_from_branch_query: {delete_from_branch_query}")
-                delete_from_branch_query_result = self.spark.sql(
-                    delete_from_branch_query
-                ).collect()
-
-                logging.info(
-                    f"result of delete_from_branch_query: {delete_from_branch_query_result}"
-                )
+        if from_tag:
+            self.disseminated_tag_df = self._gen_bronze_data_sql_from_tag()
+        else:
+            self.disseminated_tag_df = self.df_bronze
+
+        if not from_tag and dimensions is not None and len(dimensions) != 0:
+            for dimension_name, codes in dimensions.items():
+                logging.info(f"dimension_name: {dimension_name}")
+                logging.info(f"codes: {codes}")
+                if len(codes) != 0:
+                    self.disseminated_tag_df = self.disseminated_tag_df.filter(
+                        col(dimension_name).isin(codes)
+                    )
+
+        self.disseminated_tag_df.writeTo(
+            self.iceberg_tables.BRONZE_DISS_TAG.iceberg_id
+        ).createOrReplace()
+
+        logging.info(
+            f"Bronze disseminated tag table written to {self.iceberg_tables.BRONZE_DISS_TAG.iceberg_id}"
+        )
 
-        self.disseminated_tag_df = self.spark.read.option(
-            "branch", f"`diss_tag_{self.tag_name}`"
-        ).table(self.iceberg_tables.BRONZE.iceberg_id)
+        self.spark.sql(
+            f"ALTER TABLE {self.iceberg_tables.BRONZE_DISS_TAG.iceberg_id} CREATE TAG `{self.tag_name}`"
+        )
 
         disseminated_tag_df = self.disseminated_tag_df.withColumn(
             "metadata", F.to_json(col("metadata"))
@@ -481,7 +718,7 @@ class SWSBronzeIcebergSparkHelper:
         save_cache_csv(
             df=disseminated_tag_df,
             bucket=self.bucket,
-            prefix=f"{self.iceberg_tables.BRONZE.csv_prefix}_disseminated_tag",
+            prefix=f"{self.iceberg_tables.BRONZE_DISS_TAG.csv_prefix}",
             tag_name=self.tag_name,
         )
 
@@ -500,11 +737,12 @@ class SWSBronzeIcebergSparkHelper:
             private=True,
             type=TableType.ICEBERG,
             database=IcebergDatabases.BRONZE_DATABASE,
-            table=self.iceberg_tables.BRONZE.table,
-            path=self.iceberg_tables.BRONZE.path,
+            table=self.iceberg_tables.BRONZE_DISS_TAG.table,
+            path=self.iceberg_tables.BRONZE_DISS_TAG.path,
             structure={
                 "columns": self.disseminated_tag_df.schema.jsonValue()["fields"]
             },
+            pinned_columns=[*self.dim_columns_w_time, "value", *self.flag_columns],
         )
         tag = tags.add_dissemination_table(
             self.dataset_id, self.tag_name, new_iceberg_table
@@ -519,7 +757,7 @@ class SWSBronzeIcebergSparkHelper:
             private=True,
             type=TableType.CSV,
             # TODO Correct the path in the origin library
-            path=self.iceberg_tables.BRONZE.csv_path,
+            path=self.iceberg_tables.BRONZE_DISS_TAG.csv_path,
             structure={
                 "columns": self.disseminated_tag_df.schema.jsonValue()["fields"]
             },
@@ -531,3 +769,29 @@ class SWSBronzeIcebergSparkHelper:
         logging.debug(f"Tag with Added csv Table: {tag}")
 
         logging.info("Bronze Disseminated tag with selection successfully written")
+
+
+1
+frozenset({"8", "4", "2", "5", "9", "1", "7", "6", "0", "3"})
+1
+1
+2
+frozenset({"8", "4", "2", "5", "9", "1", "7", "6", "0", "3"})
+2
+1
+1
+frozenset({"8", "4", "2", "5", "9", "1", "7", "6", "0", "3"})
+1
+1
+2
+frozenset({"8", "4", "2", "5", "9", "1", "7", "6", "0", "3"})
+2
+1
+1
+frozenset({"8", "4", "2", "5", "9", "1", "7", "6", "0", "3"})
+1
+1
+1
+frozenset({"8", "4", "2", "5", "9", "1", "7", "6", "0", "3"})
+1
+1