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

sws_spark_dissemination_helper/SWSPostgresSparkReader.py

@@ -94,25 +94,37 @@ class SWSPostgresSparkReader:
 
             logging.info(f"{pg_table} read start")
 
-            # Read observations from the PostgreSQL table into a DataFrame
-            df = (
-                self.spark.read.format("jdbc")
-                .option("customSchema", custom_schema)
-                .option("dbtable", pg_table)
-                .option("partitionColumn", partition_column)
-                .option("lowerBound", min_id)
-                .option("upperBound", max_id)
-                .option("numPartitions", num_partitions)
-                .option("fetchsize", "1000")
-                .option("url", self.jdbc_url)
-                .option("user", self.jdbc_conn_properties["user"])
-                .option("password", self.jdbc_conn_properties["password"])
-                .option("driver", SPARK_POSTGRES_DRIVER)
-                .load()
-                # .repartition(1024, partition_column)
-                # .sortWithinPartitions(partition_column)
-                # .cache()
-            )
+            if min_id is None or max_id is None:
+                df = (
+                    self.spark.read.format("jdbc")
+                    .option("customSchema", custom_schema)
+                    .option("dbtable", pg_table)
+                    .option("fetchsize", "1000")
+                    .option("url", self.jdbc_url)
+                    .option("user", self.jdbc_conn_properties["user"])
+                    .option("password", self.jdbc_conn_properties["password"])
+                    .option("driver", SPARK_POSTGRES_DRIVER)
+                    .load()
+                )
+            else:
+                df = (
+                    self.spark.read.format("jdbc")
+                    .option("customSchema", custom_schema)
+                    .option("dbtable", pg_table)
+                    .option("partitionColumn", partition_column)
+                    .option("lowerBound", min_id)
+                    .option("upperBound", max_id)
+                    .option("numPartitions", num_partitions)
+                    .option("fetchsize", "1000")
+                    .option("url", self.jdbc_url)
+                    .option("user", self.jdbc_conn_properties["user"])
+                    .option("password", self.jdbc_conn_properties["password"])
+                    .option("driver", SPARK_POSTGRES_DRIVER)
+                    .load()
+                    # .repartition(1024, partition_column)
+                    # .sortWithinPartitions(partition_column)
+                    # .cache()
+                )
         else:
             df = (
                 self.spark.read.format("jdbc")
@@ -195,6 +207,7 @@ class SWSPostgresSparkReader:
             (dataset_tables.OBSERVATION_COORDINATE, "id", 10),
             (dataset_tables.METADATA, "id", 10),
             (dataset_tables.METADATA_ELEMENT, "metadata", 10),
+            (dataset_tables.TAG_OBSERVATION, "tag", 10),
         ]
         return self._import_tables(data_tables)
 
@@ -209,25 +222,30 @@ class SWSPostgresSparkReader:
             dataset_tables.METADATA_ELEMENT_TYPE,
             dataset_tables.LANGUAGE,
             dataset_tables.UNIT_OF_MEASURE,
+            dataset_tables.DATASET,
             *dataset_tables.CODELISTS,
         ]
+        logging.info(
+            f"Importing reference data tables: {[(table.postgres_id, table.iceberg_id) for table in reference_data_tables]}"
+        )
         return self._import_tables(
             [(table, None, 1) for table in reference_data_tables]
         )
 
     def import_operational_data_tables(
         self, dataset_tables: DatasetTables
-    ) -> DataFrame:
+    ) -> List[DataFrame]:
         # Define and import operational data table without partitioning
         operational_data_tables = [
             (dataset_tables.USER, None, 1),
+            (dataset_tables.TAG, None, 1),
         ]
-        return self._import_tables(operational_data_tables)[0]
+        return self._import_tables(operational_data_tables)
 
     def import_data_reference_data_operational_data(
         self, dataset_tables: DatasetTables
     ) -> Tuple[
-        Tuple[DataFrame, DataFrame, DataFrame, DataFrame],
+        Tuple[DataFrame, DataFrame, DataFrame, DataFrame, DataFrame],
         Tuple[
             DataFrame,
             DataFrame,
@@ -235,22 +253,23 @@ class SWSPostgresSparkReader:
             DataFrame,
             DataFrame,
             DataFrame,
+            DataFrame,
             List[DataFrame],
         ],
-        DataFrame,
+        Tuple[DataFrame, DataFrame],
     ]:
         # Import and organize DataFrames into the desired output structure
         data_dfs = self.import_data_tables(dataset_tables)
         reference_data_dfs = self.import_reference_data_tables(dataset_tables)
-        operational_data_df = self.import_operational_data_tables(dataset_tables)
+        operational_data_dfs = self.import_operational_data_tables(dataset_tables)
 
         return (
             tuple(data_dfs),
             (
-                *reference_data_dfs[:6],
-                reference_data_dfs[6:],
+                *reference_data_dfs[:7],
+                reference_data_dfs[7:],
             ),
-            operational_data_df,
+            tuple(operational_data_dfs),
         )
 
     def get_codelist_type_mapping(
@@ -291,13 +310,73 @@ class SWSPostgresSparkReader:
         self,
         domain_code: str,
     ) -> DataFrame:
-        return self.read_pg_table_and_check_duplicates_for_domain(
+        df = self.read_pg_table(
             pg_table=DatasetDatatables.MAPPING_CODE_CORRECTION.id,
-            table_name=DatasetDatatables.MAPPING_CODE_CORRECTION.name,
             custom_schema=DatasetDatatables.MAPPING_CODE_CORRECTION.schema,
-            domain_code=domain_code,
-            unique_columns=["old_code"],
         )
+        df.filter(
+            col("mapping_type").isNull() | (col("mapping_type") == lit(""))
+        ).transform(
+            correct_domain_filter, domain=domain_code, unique_columns=["old_code"]
+        )
+
+        return df
+
+    def get_domain_code_source_datasets_ids_dest_dataset_id(
+        self, dataset_id: str, domain_code: str = None
+    ) -> Tuple[str, List[str], str]:
+        mapping_domains_id_df = self.read_pg_table(
+            pg_table=DatasetDatatables.MAPPING_DOMAINS_ID.id,
+            custom_schema=DatasetDatatables.MAPPING_DOMAINS_ID.schema,
+        )
+
+        if domain_code is None:
+            domain_code_df = mapping_domains_id_df.filter(
+                col("sws_source_id") == lit(dataset_id)
+            ).select("domain")
+
+            if domain_code_df.count() == 0:
+                raise ValueError(
+                    f'There is no row connecting the current source dataset id ({dataset_id}) to any domain in the table "{DatasetDatatables.MAPPING_DOMAINS_ID.name}"'
+                )
+
+            if domain_code_df.count() > 1:
+                raise ValueError(
+                    f'There is more than one domain referencing the current source dataset id ({dataset_id}) in the table "{DatasetDatatables.MAPPING_DOMAINS_ID.name}", please specify the domain code you want to process in the parameters'
+                )
+
+            domain_code = domain_code_df.collect()[0][0]
+
+        source_datasets_ids = [
+            row[0]
+            for row in (
+                mapping_domains_id_df.filter(col("domain") == lit(domain_code))
+                .select("sws_source_id")
+                .collect()
+            )
+        ]
+        dest_datasets_id_df = (
+            mapping_domains_id_df.filter(col("domain") == lit(domain_code))
+            .select("sws_destination_id")
+            .distinct()
+        )
+
+        if dest_datasets_id_df.count() == 0:
+            raise ValueError(
+                f'There is no row connecting the current source dataset id and domain pair ({dataset_id}, {domain_code}) to any destination dataset id in the table "{DatasetDatatables.MAPPING_DOMAINS_ID.name}"'
+            )
+        if dest_datasets_id_df.count() > 1:
+            raise ValueError(
+                f'The source dataset id and domain pair ({dataset_id}, {domain_code}) must point only to one destination dataset in the table "{DatasetDatatables.MAPPING_DOMAINS_ID.name}"'
+            )
+
+        dest_datasets_id = dest_datasets_id_df.collect()[0][0]
+
+        logging.info(f"domain code: {domain_code}")
+        logging.info(f"source datasets ids: {source_datasets_ids}")
+        logging.info(f"dest datasets ids: {dest_datasets_id}")
+
+        return (domain_code, source_datasets_ids, dest_datasets_id)
 
     def get_dest_dataset_id(self, domain_code: str, dataset_id: str) -> Tuple[str, str]:
 
@@ -418,3 +497,45 @@ class SWSPostgresSparkReader:
                 "aggregation",
             ],
         )
+
+    def get_display_decimals_datatable(
+        self,
+        domain_code: str,
+    ) -> DataFrame:
+        df = self.read_pg_table(
+            pg_table=DatasetDatatables.DISPLAY_DECIMALS.id,
+            custom_schema=DatasetDatatables.DISPLAY_DECIMALS.schema,
+        ).filter(col("domain") == lit(domain_code))
+
+        pairs = df.select("column_1_name", "column_2_name").distinct().collect()
+
+        # If no config exists for this domain, fail early
+        if not pairs:
+            msg = (
+                f'No display-decimals configuration found for domain "{domain_code}". '
+                f'Please add an entry in table "{DatasetDatatables.DISPLAY_DECIMALS.id}".'
+            )
+            logging.error(msg)
+            # raise ValueError(msg)
+
+        # If more than one mapping exists, it's invalid
+        if len(pairs) > 1:
+            formatted_pairs = [(p["column_1_name"], p["column_2_name"]) for p in pairs]
+
+            msg = (
+                f'Invalid configuration for domain "{domain_code}". '
+                f"Expected exactly one (column_1_name, column_2_name) pair, but found {len(pairs)}: "
+                f"{formatted_pairs}. "
+                f'Please correct the table "{DatasetDatatables.DISPLAY_DECIMALS.id}".'
+            )
+
+            logging.error(
+                "Multiple display-decimals column pairs detected",
+                extra={
+                    "domain": domain_code,
+                    "pairs_found": formatted_pairs,
+                },
+            )
+            raise ValueError(msg)
+
+        return df

sws_spark_dissemination_helper/SWSSilverIcebergSparkHelper.py

@@ -10,7 +10,7 @@ from pyspark.sql.window import Window
 from sws_api_client import Tags
 from sws_api_client.tags import BaseDisseminatedTagTable, TableLayer, TableType
 
-from .constants import IcebergDatabases, IcebergTables
+from .constants import IcebergDatabases, IcebergTables, DatasetDatatables
 from .SWSPostgresSparkReader import SWSPostgresSparkReader
 from .utils import (
     get_or_create_tag,
@@ -103,7 +103,7 @@ class SWSSilverIcebergSparkHelper:
         # The diss_flag column is needed to initialize the condition expression
         # The note column will contain the eventual reasons why diss_flag has been set to false
         return df.withColumn("diss_flag", lit(True)).withColumn(
-            "note", lit([]).cast(ArrayType(StringType()))
+            "diss_note", lit([]).cast(ArrayType(StringType()))
         )
 
     def read_bronze_data(self) -> DataFrame:
@@ -111,6 +111,11 @@ class SWSSilverIcebergSparkHelper:
             self.iceberg_tables.BRONZE.iceberg_id
         )
 
+    def read_bronze_diss_tag_data(self) -> DataFrame:
+        return self.spark.read.option("tag", self.tag_name).table(
+            self.iceberg_tables.BRONZE_DISS_TAG.iceberg_id
+        )
+
     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."""
         dim_columns_w_time = [
@@ -158,6 +163,99 @@ class SWSSilverIcebergSparkHelper:
 
         logging.info(f"Checking time validity for {col_name} of type {col_type}")
 
+        if col_type == "area":
+            logging.info(
+                f'Changing start and end year according to "{DatasetDatatables.MAPPING_CODE_CORRECTION.name}"'
+            )
+            df_start_year_correction = self.df_mapping_code_correction.filter(
+                col("var_type") == lit("start_year")
+            )
+            df_end_year_correction = self.df_mapping_code_correction.filter(
+                col("var_type") == lit("end_year")
+            )
+
+            original_col_order = df.columns
+            cols_to_select = df.columns
+            col_name_lower = col_name.lower()
+            cols_to_select = [
+                column
+                for column in cols_to_select
+                if column.lower()
+                not in (
+                    "diss_note",
+                    f"{col_name_lower}_start_date",
+                    f"{col_name_lower}_end_date",
+                )
+            ]
+
+            df = (
+                df.alias("d")
+                .join(
+                    F.broadcast(df_start_year_correction).alias("sy"),
+                    on=col(f"d.{col_name}") == col("sy.mapping_type"),
+                    how="left",
+                )
+                .join(
+                    F.broadcast(df_end_year_correction).alias("ey"),
+                    on=col(f"d.{col_name}") == col("ey.mapping_type"),
+                    how="left",
+                )
+                .withColumn("valid_new_start_year", col("sy.new_code").isNotNull())
+                .withColumn("valid_new_end_year", col("ey.new_code").isNotNull())
+                .withColumn(
+                    "new_diss_note",
+                    F.when(
+                        col("valid_new_start_year"),
+                        F.array_append(
+                            col("d.diss_note"),
+                            F.concat(
+                                col("sy.note"),
+                                lit(" from "),
+                                col("sy.old_code"),
+                                lit(" to "),
+                                col("sy.new_code"),
+                            ),
+                        ),
+                    ).otherwise(col("d.diss_note")),
+                )
+                .withColumn(
+                    "new_diss_note",
+                    F.when(
+                        col("valid_new_end_year"),
+                        F.array_append(
+                            col("new_diss_note"),
+                            F.concat(
+                                col("ey.note"),
+                                lit(" from "),
+                                col("ey.old_code"),
+                                lit(" to "),
+                                col("ey.new_code"),
+                            ),
+                        ),
+                    ).otherwise(col("new_diss_note")),
+                )
+                .withColumn(
+                    f"new_{col_name}_start_date",
+                    F.when(
+                        col("valid_new_start_year"), F.to_date(col("sy.new_code"))
+                    ).otherwise(col(f"d.{col_name}_start_date")),
+                )
+                .withColumn(
+                    f"new_{col_name}_end_date",
+                    F.when(
+                        col("valid_new_end_year"),
+                        F.to_date(F.concat(col("ey.new_code"), lit("-12-31"))),
+                    ).otherwise(col(f"d.{col_name}_end_date")),
+                )
+                .select(
+                    *cols_to_select,
+                    col("new_diss_note").alias("diss_note"),
+                    col(f"new_{col_name}_start_date").alias(f"{col_name}_start_date"),
+                    col(f"new_{col_name}_end_date").alias(f"{col_name}_end_date"),
+                )
+                .select(*original_col_order)
+            )
+
         # Iterate through columns and build conditions dynamically
         start_date_condition = col(f"{col_name}_start_date").isNull() | (
             col(f"{col_name}_start_date") <= col(f"{self.time_column}_start_date")
@@ -172,15 +270,15 @@ class SWSSilverIcebergSparkHelper:
                 start_date_condition & end_date_condition,
             )
             .withColumn("diss_flag", col("diss_flag") & col("condition_result"))
-            # In case the condition is satisfied update diss_flag accordingly and append a note indicating the reason for the observation exclusion from the dissemination
+            # In case the condition is satisfied update diss_flag accordingly and append a diss_note indicating the reason for the observation exclusion from the dissemination
             .withColumn(
-                "note",
+                "diss_note",
                 F.when(
                     ~col("condition_result"),
                     F.array_append(
-                        col("note"), lit(f"{col_type} out of time validity range")
+                        col("diss_note"), lit(f"{col_type} out of time validity range")
                     ),
-                ).otherwise(col("note")),
+                ).otherwise(col("diss_note")),
             )
             .drop("condition_result")
         )
@@ -292,7 +390,7 @@ class SWSSilverIcebergSparkHelper:
             col_name (str): The DataFrame column name on which to apply the filter
 
         Returns:
-            DataFrame: The DataFrame with updated `diss_flag` and `note` columns based on the check outcome
+            DataFrame: The DataFrame with updated `diss_flag` and `diss_note` columns based on the check outcome
         """
 
         # Remove the duplicates that may be in the tables
@@ -330,14 +428,14 @@ class SWSSilverIcebergSparkHelper:
                 col("diss_flag") & col("condition_result"),
             )
             .withColumn(
-                "note",
+                "diss_note",
                 F.when(
                     ~col("condition_result"),
                     F.array_append(
-                        col("note"),
+                        col("diss_note"),
                         lit(f"{col_type} not disseminated for this domain"),
                     ),
-                ).otherwise(col("note")),
+                ).otherwise(col("diss_note")),
             )
             .drop("condition_result")
         )
@@ -424,16 +522,16 @@ class SWSSilverIcebergSparkHelper:
                             col("diss_flag") & col("condition_result"),
                         )
                         .withColumn(
-                            "note",
+                            "diss_note",
                             F.when(
                                 ~col("condition_result"),
                                 F.array_append(
-                                    col("note"),
+                                    col("diss_note"),
                                     lit(
                                         f"not disseminated according to exception with note: {row_exception['note']}"
                                     ),
                                 ),
-                            ).otherwise(col("note")),
+                            ).otherwise(col("diss_note")),
                         )
                         .drop("condition_result")
                     )
@@ -518,7 +616,7 @@ class SWSSilverIcebergSparkHelper:
 
         df = (
             df.withColumn("metadata", F.to_json(col("metadata")))
-            .withColumn("note", F.to_json(col("note")))
+            .withColumn("diss_note", F.to_json(col("diss_note")))
             .coalesce(1)
         )
 
@@ -547,6 +645,13 @@ class SWSSilverIcebergSparkHelper:
             table=self.iceberg_tables.SILVER.table,
             path=self.iceberg_tables.SILVER.path,
             structure={"columns": df.schema.jsonValue()["fields"]},
+            pinned_columns=[
+                *self.dim_columns_w_time,
+                "value",
+                *self.flag_columns,
+                "diss_flag",
+                "diss_note",
+            ],
         )
         tag = upsert_disseminated_table(
             sws_tags=tags,

sws_spark_dissemination_helper/__init__.py

@@ -2,3 +2,4 @@ from .SWSPostgresSparkReader import SWSPostgresSparkReader
 from .SWSBronzeIcebergSparkHelper import SWSBronzeIcebergSparkHelper
 from .SWSSilverIcebergSparkHelper import SWSSilverIcebergSparkHelper
 from .SWSGoldIcebergSparkHelper import SWSGoldIcebergSparkHelper
+from .SWSEasyIcebergSparkHelper import SWSEasyIcebergSparkHelper