feat: implement discovery-first global ID mapping to eliminate redundant registry I/O; update unit tests accordingly

Author: BLMgithub

data_pipeline/assembly/assembly_executor.py

@@ -147,6 +147,10 @@ def orchestrate_event_assembly(run_context: RunContext, report: Dict) -> bool:
         log_error(f"Unexpected error processing event assembly: {e}", report)
 
     finally:
+        if "lf_derived" in locals():
+            del lf_derived  # type: ignore
+        if "lf_freezed" in locals():
+            del lf_freezed
         force_gc()
 
     return True
@@ -182,9 +186,6 @@ def orchestrate_dimension_refs(run_context: RunContext, report: Dict) -> bool:
         report[table] = {"dim_reference": False, "export": False}
         tracker = report[table]
 
-        lf_raw = None
-        df_dim = None
-
         try:
             # Switch between local and gcp IO
             if run_context.bq_project_id == "PROJECT_ID_NOT_DETECTED":
@@ -246,10 +247,10 @@ def orchestrate_dimension_refs(run_context: RunContext, report: Dict) -> bool:
             return False
 
         finally:
-            if lf_raw is not None:
-                del lf_raw
-            if df_dim is not None:
-                del df_dim
+            if "lf_raw" in locals():
+                del lf_raw  # type: ignore
+            if "df_dim" in locals():
+                del df_dim  # type: ignore
             gc.collect()
 
     return True

data_pipeline/contract/contract_executor.py

@@ -2,44 +2,45 @@
 # Contract Stage Executor
 # =============================================================================
 
+import polars as pl
 from data_pipeline.shared.run_context import RunContext
 from data_pipeline.shared.loader_exporter import load_single_delta, export_file
+from data_pipeline.assembly.assembly_executor import force_gc
 from data_pipeline.shared.table_configs import TABLE_CONFIG
 from data_pipeline.contract.registry import ROLE_STEPS
-from data_pipeline.contract.id_registrar import id_mapping, ID_COLUMNS_TO_MAP
-from pathlib import Path
+from data_pipeline.contract.id_registrar import ID_ENTITY_MAP
 
 
 def apply_contract(
     run_context: RunContext,
     table_name: str,
+    master_mappings: dict[str, pl.LazyFrame],
     invalid_order_ids: set | None = None,
     valid_order_ids: set | None = None,
 ) -> tuple[dict, set, set]:
     """
-    Main entry point for the Raw-to-Contracted Stage.
+    Orchestrates the Raw-to-Contracted transformation for a specific logical table.
 
     Workflow:
     1. Resolve: Identifies table metadata (role, schema, keys) from the central registry.
     2. Hydrate: Fetches the raw snapshot from the lake's snapshot zone.
     3. Delegate: Iteratively applies atomic logic rules (Deduplication, Chronology, Null-checks).
     4. Validate: Executes 'enforce_schema' as the terminal structural gate.
-    5. Promote: Persists the contract-compliant dataset to the Silver (contracted) zone.
+    5. Map: Joins against pre-calculated Discovery mappings to enrich UUIDs with UInt32 integer IDs.
+    6. Promote: Persists the contract-compliant dataset to the Silver (contracted) zone.
 
     Operational Guarantees:
     - Subtractive Only: Exclusively filters rows or casts types; never mutates business values.
     - Referential Safety: Propagates invalidated keys across table boundaries to ensure consistent pruning.
-    - Structural Finality: Guarantees output parity with the ASSEMBLE_SCHEMA specification.
+    - Structural Finality: Guarantees output parity with the Silver layer specification, including required integer IDs.
 
     Side Effects:
     - Persists a Parquet artifact to the contracted directory.
-    - Updates newly invalidated 'order_id' sets for downstream cross-table pruning.
+    - Updates invalidated 'order_id' sets for downstream cross-table pruning.
 
     Failure Behavior:
     - Traps logic-step exceptions; logs errors to the report and halts the current table's processing.
-
-    Returns:
-        tuple: (Stage Report Dict, Newly Invalidated IDs Set, Validated Order IDs Set)
+    - Crashes if ID mapping joins fail to prevent downstream schema corruption.
     """
 
     report = {
@@ -69,7 +70,6 @@ def apply_contract(
     if table_name not in TABLE_CONFIG:
         report["status"] = "failed"
         report["errors"].append(f"Unknown table: {table_name}")
-
         return report, invalid_ids, valid_ids
 
     base_path = run_context.raw_snapshot_path
@@ -89,32 +89,26 @@ def apply_contract(
     role = config["role"]
 
     for step in ROLE_STEPS[role]:
-
         contract = step["contract"]
-        args = []
-
-        if "non_nullable" in step["args"]:
-            args.append(non_nullable)
 
-        if "invalid_order_ids" in step["args"]:
-            args.append(invalid_order_ids)
+        args = [
+            non_nullable if "non_nullable" in step["args"] else None,
+            invalid_order_ids if "invalid_order_ids" in step["args"] else None,
+            valid_order_ids if "valid_order_ids" in step["args"] else None,
+        ]
 
-        if "valid_order_ids" in step["args"]:
-            args.append(valid_order_ids)
+        # Remove args not needed for the current registry loop
+        args = [arg for arg in args if arg is not None]
 
         if "required_column" in step["args"]:
-            args.append(required_column)
-            args.append(dtypes)
+            args.extend([required_column, dtypes])
 
         try:
-
             if step["return_invalid_ids"]:
                 df, removed, new_invalid = contract(df, *args)
                 invalid_ids |= new_invalid
-
             else:
                 df, removed = contract(df, *args)
-
             report[step["metric"]] += removed
 
         except Exception as e:
@@ -129,33 +123,31 @@ def apply_contract(
     if table_name == "df_orders":
         valid_ids = set(df.get_column("order_id"))
 
-    temp_path = run_context.contracted_path / "id_mapping"
-    storage_mapping_str = run_context.storage_mapping_path
-
-    # Switch between local and gcp path
-    storage_mapping = (
-        storage_mapping_str
-        if storage_mapping_str.startswith("gs://")
-        else Path(storage_mapping_str)
-    )
-
-    df = id_mapping(
-        df=df,
-        table_name=table_name,
-        mapping_dict=ID_COLUMNS_TO_MAP,
-        runtime_dir=temp_path,
-        destination=storage_mapping,
-        run_id=run_context.run_id,
-    )
+    df_lf = df.lazy()
 
-    output_path = run_context.contracted_path / f"{filename}.parquet"
+    try:
+        # Attach mapped integer in dataframe
+        for entity_col, tables in ID_ENTITY_MAP.items():
+            if table_name in tables and entity_col in master_mappings:
+                df_lf = df_lf.join(
+                    master_mappings[entity_col], on=entity_col, how="left"
+                )
+
+    # Force to fail before corrupting downstream
+    except Exception as e:
+        raise RuntimeError(f"Mapping Uint32 to UUIDs Failed: {e}") from e
 
-    if not export_file(df, output_path):
+    output_path = run_context.contracted_path / f"{filename}.parquet"
+    if not export_file(df_lf.collect(), output_path):
         report["status"] = "failed"
         report["errors"].append("Export failed")
-
         return report, invalid_ids, valid_ids
 
-    report["status"] = "success"
+    if "df" in locals():
+        del df
+    if "df_lf" in locals():
+        del df_lf
+    force_gc()
 
+    report["status"] = "success"
     return report, invalid_ids, valid_ids

data_pipeline/contract/contract_logic.py

@@ -6,7 +6,6 @@
 from typing import List
 from data_pipeline.shared.table_configs import REQUIRED_TIMESTAMPS, TIMESTAMP_FORMATS
 
-
 # ------------------------------------------------------------
 # CONTRACT LOGICS
 # ------------------------------------------------------------
@@ -68,6 +67,7 @@ def remove_unparsable_timestamps(df: pl.DataFrame) -> tuple[pl.DataFrame, int, s
     exprs = []
     for col in REQUIRED_TIMESTAMPS:
         if col in df.columns:
+
             if df.schema[col] == pl.String:
                 fmt = TIMESTAMP_FORMATS.get(col)
                 exprs.append(
@@ -255,6 +255,7 @@ def enforce_schema(
         target_dtype = dtypes.get(col)
 
         if target_dtype == pl.Datetime:
+
             if df.schema[col] == pl.String:
                 fmt = TIMESTAMP_FORMATS.get(col)
                 exprs.append(pl.col(col).str.to_datetime(format=fmt, strict=False))