contractforge-databricks 0.1.0__py3-none-any.whl

contractforge_databricks/parity/catalog.py

@@ -0,0 +1,59 @@
+"""Query helpers for the Databricks parity catalog."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.parity import WriteEngineParityScenario
+from contractforge_databricks.parity.scenarios import WRITE_ENGINE_PARITY_SCENARIOS
+
+
+def list_write_engine_parity_scenarios() -> list[str]:
+    return sorted(scenario.scenario_id for scenario in WRITE_ENGINE_PARITY_SCENARIOS)
+
+
+def get_write_engine_parity_scenario(scenario_id: str) -> WriteEngineParityScenario:
+    for scenario in WRITE_ENGINE_PARITY_SCENARIOS:
+        if scenario.scenario_id == scenario_id:
+            return scenario
+    raise ValueError(
+        f"Write-engine parity scenario not found: {scenario_id}. "
+        f"Valid scenarios: {list_write_engine_parity_scenarios()}"
+    )
+
+
+def scenarios_for_engine(engine: str) -> list[WriteEngineParityScenario]:
+    return [scenario for scenario in WRITE_ENGINE_PARITY_SCENARIOS if scenario.candidate_engine == engine]
+
+
+def scenarios_for_mode(mode: str) -> list[WriteEngineParityScenario]:
+    return [scenario for scenario in WRITE_ENGINE_PARITY_SCENARIOS if scenario.write_mode == mode]
+
+
+def build_write_engine_parity_plan(
+    *,
+    engine: str | None = None,
+    mode: str | None = None,
+    runtime: str | None = None,
+) -> dict[str, Any]:
+    scenarios = list(WRITE_ENGINE_PARITY_SCENARIOS)
+    if engine:
+        scenarios = [scenario for scenario in scenarios if scenario.candidate_engine == engine]
+    if mode:
+        scenarios = [scenario for scenario in scenarios if scenario.write_mode == mode]
+    if runtime:
+        scenarios = [scenario for scenario in scenarios if runtime in scenario.runtime_targets]
+
+    expectation_counts: dict[str, int] = {}
+    for scenario in scenarios:
+        expectation_counts[scenario.expectation] = expectation_counts.get(scenario.expectation, 0) + 1
+
+    return {
+        "kind": "write_engine_parity_plan",
+        "engine": engine,
+        "mode": mode,
+        "runtime": runtime,
+        "scenario_count": len(scenarios),
+        "expectation_counts": expectation_counts,
+        "scenarios": [scenario.as_dict() for scenario in scenarios],
+    }

contractforge_databricks/parity/models.py

@@ -0,0 +1,7 @@
+"""Compatibility exports for platform-neutral parity catalog models."""
+
+from contractforge_core.parity import ParityExpectation, ParityMetricExpectation, WriteEngineParityScenario
+
+RuntimeTarget = str
+
+__all__ = ["ParityExpectation", "ParityMetricExpectation", "RuntimeTarget", "WriteEngineParityScenario"]

contractforge_databricks/parity/scenarios.py

@@ -0,0 +1,111 @@
+"""Official Databricks write-engine parity scenarios."""
+
+from __future__ import annotations
+
+from contractforge_core.parity import ParityMetricExpectation, WriteEngineParityScenario
+
+CLASSIC_AND_SERVERLESS = ("databricks_classic", "databricks_serverless")
+
+WRITE_ENGINE_PARITY_SCENARIOS: tuple[WriteEngineParityScenario, ...] = (
+    WriteEngineParityScenario(
+        scenario_id="scd1_sql_merge_insert_update",
+        title="SCD1 SQL MERGE preserves insert/update semantics",
+        write_mode="scd1_upsert",
+        candidate_engine="databricks_sql_merge",
+        expectation="must_match",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("databricks_runtime", "unity_catalog_table", "sql_merge"),
+        required_contract_fields=("merge_keys", "transform.deduplicate"),
+        expected_semantics=("one current row per merge key", "changed rows update", "new keys insert"),
+        metric_expectations=(
+            ParityMetricExpectation("rows_inserted", "must match ContractForge Delta"),
+            ParityMetricExpectation("rows_updated", "must match ContractForge Delta"),
+        ),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="scd1_sql_merge_duplicate_keys",
+        title="SCD1 SQL MERGE handles duplicate keys deterministically",
+        write_mode="scd1_upsert",
+        candidate_engine="databricks_sql_merge",
+        expectation="must_match",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("databricks_runtime", "unity_catalog_table", "sql_merge"),
+        required_contract_fields=("merge_keys", "transform.deduplicate"),
+        expected_semantics=("source row selected per key must match ContractForge deduplication",),
+        metric_expectations=(ParityMetricExpectation("rows_affected", "must match after deduplication"),),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="scd1_sql_merge_null_keys",
+        title="SCD1 SQL MERGE rejects or quarantines null merge keys consistently",
+        write_mode="scd1_upsert",
+        candidate_engine="databricks_sql_merge",
+        expectation="must_match",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("databricks_runtime", "unity_catalog_table", "sql_merge"),
+        required_contract_fields=("merge_keys", "quality_rules.not_null"),
+        expected_semantics=("null merge-key rows follow declared quality behavior",),
+        metric_expectations=(ParityMetricExpectation("rows_quarantined", "must match when quarantine is configured"),),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="scd2_auto_cdc_history_lifecycle",
+        title="Lakeflow AUTO CDC SCD2 preserves history lifecycle",
+        write_mode="scd2_historical",
+        candidate_engine="databricks_lakeflow_auto_cdc",
+        expectation="must_match",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("lakeflow_auto_cdc", "unity_catalog_table"),
+        required_contract_fields=("merge_keys", "sequence_by", "scd2_change_columns"),
+        expected_semantics=("changed columns expire current row", "new current version is inserted"),
+        metric_expectations=(ParityMetricExpectation("rows_expired", "normalized metric must match"),),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="scd2_auto_cdc_late_arriving",
+        title="Lakeflow AUTO CDC SCD2 handles late-arriving changes explicitly",
+        write_mode="scd2_historical",
+        candidate_engine="databricks_lakeflow_auto_cdc",
+        expectation="must_match",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("lakeflow_auto_cdc", "unity_catalog_table"),
+        required_contract_fields=("merge_keys", "sequence_by", "scd2_late_arriving_policy"),
+        expected_semantics=("late-arriving records must not corrupt current row",),
+        metric_expectations=(ParityMetricExpectation("rows_expired", "must match or document intentional difference"),),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="scd2_auto_cdc_delete_semantics",
+        title="Lakeflow AUTO CDC SCD2 delete predicates preserve declared delete behavior",
+        write_mode="scd2_historical",
+        candidate_engine="databricks_lakeflow_auto_cdc",
+        expectation="must_match",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("lakeflow_auto_cdc", "unity_catalog_table"),
+        required_contract_fields=("merge_keys", "sequence_by", "apply_as_deletes"),
+        expected_semantics=("delete predicates must be explicit and reviewed",),
+        metric_expectations=(ParityMetricExpectation("rows_deleted", "must match declared delete behavior"),),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="hash_diff_auto_cdc_non_equivalence",
+        title="Hash-diff append is not equivalent to Lakeflow SCD type 1",
+        write_mode="scd1_hash_diff",
+        candidate_engine="databricks_lakeflow_auto_cdc",
+        expectation="unsupported",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("lakeflow_auto_cdc",),
+        required_contract_fields=("hash_keys",),
+        expected_semantics=("hash diff appends changed versions; Lakeflow SCD1 stores current state",),
+        metric_expectations=(ParityMetricExpectation("rows_inserted", "must not be compared"),),
+        blockers_to_record=("scd1_hash_diff is append-only version capture, not Lakeflow SCD type 1.",),
+    ),
+    WriteEngineParityScenario(
+        scenario_id="snapshot_soft_delete_auto_cdc_difference",
+        title="Snapshot soft delete differs from AUTO CDC snapshot deletes",
+        write_mode="snapshot_soft_delete",
+        candidate_engine="databricks_lakeflow_auto_cdc",
+        expectation="intentional_difference",
+        runtime_targets=CLASSIC_AND_SERVERLESS,
+        required_capabilities=("lakeflow_auto_cdc",),
+        required_contract_fields=("merge_keys",),
+        expected_semantics=("ContractForge marks missing keys inactive; snapshot CDC may delete missing keys",),
+        metric_expectations=(ParityMetricExpectation("rows_deleted", "expected to differ"),),
+        blockers_to_record=("AUTO CDC snapshot deletes are not inactive-marker soft deletes.",),
+    ),
+)

contractforge_databricks/partitioning/__init__.py

@@ -0,0 +1,3 @@
+from contractforge_databricks.partitioning.predicates import render_partition_in_predicate, render_replace_where
+
+__all__ = ["render_partition_in_predicate", "render_replace_where"]

contractforge_databricks/partitioning/predicates.py

@@ -0,0 +1,28 @@
+"""Databricks partition predicate rendering."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+from contractforge_core.partitioning import distinct_partition_values
+from contractforge_databricks.sql import quote_identifier, sql_string
+
+
+def render_partition_in_predicate(column: str, values: Iterable[Any], *, max_values: int = 1000) -> str:
+    distinct = distinct_partition_values(values, max_values=max_values)
+    quoted = quote_identifier(column)
+    non_null = [value for value in distinct if value is not None]
+    predicates = []
+    if non_null:
+        literals = ", ".join(sql_string(value) for value in non_null)
+        predicates.append(f"{quoted} IN ({literals})")
+    if any(value is None for value in distinct):
+        predicates.append(f"{quoted} IS NULL")
+    return " OR ".join(predicates)
+
+
+def render_replace_where(column: str, value: Any) -> str:
+    if value is None:
+        return f"{quote_identifier(column)} IS NULL"
+    return f"{quote_identifier(column)} = {sql_string(value)}"

contractforge_databricks/preparation/__init__.py

@@ -0,0 +1,47 @@
+from contractforge_core.preparation import HashDiffStageSpec, SCD2StageSpec, SnapshotStageSpec
+from contractforge_databricks.preparation.hashing import (
+    HASH_DELIMITER,
+    HASH_NULL_SENTINEL,
+    ROW_HASH_COLUMN,
+    render_row_hash_expression,
+)
+from contractforge_databricks.preparation.encoding import apply_encoding_fix
+from contractforge_databricks.preparation.shape import apply_shape
+from contractforge_databricks.preparation.deduplicate import apply_transform_deduplicate
+from contractforge_databricks.preparation.pyspark import (
+    apply_transform,
+    apply_contract_preparation,
+    apply_transform_cast,
+    apply_transform_derive,
+    apply_transform_standardize,
+)
+from contractforge_databricks.preparation.pyspark_staging import (
+    apply_write_staging,
+    prepare_hash_diff_stage,
+    prepare_scd2_stage,
+    prepare_snapshot_stage,
+    with_row_hash,
+)
+
+__all__ = [
+    "HashDiffStageSpec",
+    "HASH_DELIMITER",
+    "HASH_NULL_SENTINEL",
+    "ROW_HASH_COLUMN",
+    "SCD2StageSpec",
+    "SnapshotStageSpec",
+    "apply_encoding_fix",
+    "apply_shape",
+    "apply_contract_preparation",
+    "apply_transform",
+    "apply_write_staging",
+    "apply_transform_cast",
+    "apply_transform_deduplicate",
+    "apply_transform_derive",
+    "apply_transform_standardize",
+    "prepare_hash_diff_stage",
+    "prepare_scd2_stage",
+    "prepare_snapshot_stage",
+    "render_row_hash_expression",
+    "with_row_hash",
+]

contractforge_databricks/preparation/deduplicate.py

@@ -0,0 +1,87 @@
+"""PySpark deduplication helpers for portable transform intent."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+
+def apply_transform_deduplicate(df: Any, deduplicate: object) -> Any:
+    if not isinstance(deduplicate, dict):
+        return df
+    from pyspark.sql import Window, functions as F
+
+    keys = deduplicate.get("keys")
+    key_columns = [str(keys)] if isinstance(keys, str) else [str(key) for key in keys or ()]
+    if not key_columns:
+        raise ValueError("transform.deduplicate.keys is required")
+    _validate_columns(df, {column: True for column in key_columns}, "transform.deduplicate.keys")
+    window = Window.partitionBy(*key_columns).orderBy(*_deduplicate_order_columns(deduplicate.get("order_by"), F))
+    return df.withColumn("__cf_row_number", F.row_number().over(window)).filter(F.col("__cf_row_number") == 1).drop(
+        "__cf_row_number"
+    )
+
+
+def _deduplicate_order_columns(order_by: object, functions: Any) -> list[Any]:
+    if isinstance(order_by, str):
+        return _deduplicate_order_columns_from_string(order_by, functions)
+    order_columns = []
+    for item in order_by or ():
+        if not isinstance(item, dict):
+            continue
+        order_columns.append(
+            _deduplicate_order_column(
+                functions.col(str(item["column"])),
+                direction=str(item.get("direction", "desc")).lower(),
+                nulls=str(item.get("nulls") or "").lower(),
+            )
+        )
+    if not order_columns:
+        raise ValueError("transform.deduplicate.order_by is required")
+    return order_columns
+
+
+def _deduplicate_order_columns_from_string(order_by: str, functions: Any) -> list[Any]:
+    order_columns = []
+    for clause in (item.strip() for item in order_by.split(",")):
+        if not clause:
+            continue
+        parsed = re.match(
+            r"^`?(?P<column>[A-Za-z_][A-Za-z0-9_]*)`?(?:\s+(?P<direction>ASC|DESC))?(?:\s+NULLS\s+(?P<nulls>FIRST|LAST))?$",
+            clause,
+            flags=re.IGNORECASE,
+        )
+        if parsed is None:
+            order_columns.append(functions.expr(clause))
+            continue
+        order_columns.append(
+            _deduplicate_order_column(
+                functions.col(parsed.group("column")),
+                direction=(parsed.group("direction") or "desc").lower(),
+                nulls=(parsed.group("nulls") or "").lower(),
+            )
+        )
+    if not order_columns:
+        raise ValueError("transform.deduplicate.order_by is required")
+    return order_columns
+
+
+def _deduplicate_order_column(column: Any, *, direction: str, nulls: str) -> Any:
+    if direction == "asc" and nulls == "first":
+        return column.asc_nulls_first()
+    if direction == "asc" and nulls == "last":
+        return column.asc_nulls_last()
+    if direction == "asc":
+        return column.asc()
+    if nulls == "first":
+        return column.desc_nulls_first()
+    if nulls == "last":
+        return column.desc_nulls_last()
+    return column.desc()
+
+
+def _validate_columns(df: Any, columns: dict[str, Any], context: str) -> None:
+    available = set(getattr(df, "columns", ()) or ())
+    missing = sorted(str(column) for column in columns if str(column) not in available)
+    if missing:
+        raise ValueError(f"{context} references missing columns: {missing}")

contractforge_databricks/preparation/encoding.py

@@ -0,0 +1,37 @@
+"""Databricks preparation helpers for adapter-owned encoding fixes."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+from contractforge_databricks.contract_extensions import databricks_extensions
+
+
+def apply_encoding_fix(df: Any, contract: Any) -> Any:
+    extensions = databricks_extensions(contract)
+    if not extensions.get("fix_encoding"):
+        return df
+    functions = importlib.import_module("pyspark.sql").functions
+    encoding = str(extensions.get("encoding") or "utf-8")
+    string_columns = _string_columns(df)
+    for column in _string_tuple(extensions.get("encoding_columns")) or string_columns:
+        if column in string_columns:
+            df = df.withColumn(column, functions.decode(functions.col(column).cast("binary"), encoding))
+    return df
+
+
+def _string_columns(df: Any) -> tuple[str, ...]:
+    return tuple(
+        field.name
+        for field in getattr(getattr(df, "schema", None), "fields", ()) or ()
+        if field.dataType.typeName() == "string"
+    )
+
+
+def _string_tuple(value: object) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        return tuple(part.strip() for part in value.split(",") if part.strip())
+    return tuple(str(part) for part in value or ())

contractforge_databricks/preparation/hashing.py

@@ -0,0 +1,18 @@
+"""Hash expression helpers for Databricks staging."""
+
+from __future__ import annotations
+
+from contractforge_core.preparation import HASH_DELIMITER, HASH_NULL_SENTINEL
+from contractforge_databricks.sql import quote_identifier
+
+ROW_HASH_COLUMN = "row_hash"
+
+
+def render_row_hash_expression(columns: tuple[str, ...], *, exclude: tuple[str, ...] = ()) -> str:
+    included = tuple(column for column in columns if column not in set(exclude))
+    if not included:
+        raise ValueError("row hash requires at least one included column")
+    payload = ", ".join(
+        f"coalesce(cast({quote_identifier(column)} as string), '{HASH_NULL_SENTINEL}')" for column in included
+    )
+    return f"sha2(concat_ws('{HASH_DELIMITER}', {payload}), 256)"

contractforge_databricks/preparation/pyspark.py

@@ -0,0 +1,178 @@
+"""Optional PySpark staging helpers.
+
+Imports stay inside functions so the package can be imported without PySpark.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.config import CONTROL_COLUMNS
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.preparation.deduplicate import apply_transform_deduplicate
+from contractforge_databricks.preparation.encoding import apply_encoding_fix
+
+
+def create_or_replace_temp_view(df: Any, name: str) -> str:
+    if not name or not name.strip():
+        raise ValueError("temp view name must not be empty")
+    df.createOrReplaceTempView(name)
+    return name
+
+
+def apply_transform(df: Any, transform: dict[str, Any] | None) -> Any:
+    """Apply portable transform intent with PySpark DataFrame operations."""
+
+    if not transform:
+        return df
+    df = apply_transform_cast(df, _dict(transform.get("cast")))
+    df = apply_transform_standardize(df, _dict(transform.get("standardize")))
+    df = apply_transform_derive(df, _dict(transform.get("derive")))
+    df = apply_transform_composite_keys(df, _dict(transform.get("composite_keys")))
+    return apply_transform_deduplicate(df, transform.get("deduplicate"))
+
+
+def apply_contract_preparation(
+    df: Any,
+    contract: SemanticContract,
+    *,
+    watermark_column: str | None = None,
+    watermark_previous: str | None = None,
+) -> Any:
+    """Apply portable pre-write preparation declared in the core contract."""
+
+    metadata = _contract_metadata(contract)
+    select_columns = _string_list(metadata.get("select_columns"))
+    if select_columns:
+        _validate_columns(df, {column: True for column in select_columns}, "select_columns")
+        df = df.select(*select_columns)
+    column_mapping = _dict(metadata.get("column_mapping"))
+    if column_mapping:
+        _validate_column_mapping(df, column_mapping)
+        for source_col, target_col in column_mapping.items():
+            df = df.withColumnRenamed(str(source_col), str(target_col))
+    if contract.shape:
+        from contractforge_databricks.preparation.shape import apply_shape
+
+        df = apply_shape(df, contract.shape.raw, layer=contract.target.layer)
+    transform = contract.transform.raw if contract.transform else {}
+    df = apply_transform_cast(df, _dict(transform.get("cast")))
+    df = apply_transform_standardize(df, _dict(transform.get("standardize")))
+    df = apply_transform_derive(df, _dict(transform.get("derive")))
+    filter_expression = metadata.get("filter_expression")
+    if filter_expression:
+        from pyspark.sql import functions as F
+
+        df = df.where(F.expr(str(filter_expression)))
+    df = apply_transform_composite_keys(df, _dict(transform.get("composite_keys")))
+    df = _apply_watermark_filter(df, watermark_column, watermark_previous)
+    df = apply_transform_deduplicate(df, transform.get("deduplicate"))
+    return apply_encoding_fix(df, contract)
+
+
+def apply_transform_cast(df: Any, casts: dict[str, Any]) -> Any:
+    if not casts:
+        return df
+    from pyspark.sql import functions as F
+
+    _validate_columns(df, casts, "transform.cast")
+    for column_name, data_type in casts.items():
+        df = df.withColumn(str(column_name), F.col(str(column_name)).cast(str(data_type)))
+    return df
+
+
+def apply_transform_derive(df: Any, expressions: dict[str, Any]) -> Any:
+    if not expressions:
+        return df
+    from pyspark.sql import functions as F
+
+    for column_name, expression in expressions.items():
+        df = df.withColumn(str(column_name), F.expr(str(expression)))
+    return df
+
+
+def apply_transform_composite_keys(df: Any, composite_keys: dict[str, Any]) -> Any:
+    if not composite_keys:
+        return df
+    from pyspark.sql import functions as F
+
+    for key_name, source_columns in composite_keys.items():
+        columns = [source_columns] if isinstance(source_columns, str) else list(source_columns or ())
+        _validate_columns(df, {str(column): True for column in columns}, f"transform.composite_keys.{key_name}")
+        parts = [F.coalesce(F.col(str(column)).cast("string"), F.lit("")) for column in columns]
+        df = df.withColumn(str(key_name), F.concat_ws("|", *parts))
+    return df
+
+
+def apply_transform_standardize(df: Any, standardize: dict[str, Any]) -> Any:
+    if not standardize:
+        return df
+    from pyspark.sql import functions as F
+
+    _validate_columns(df, standardize, "transform.standardize")
+    for column_name, config in standardize.items():
+        column = F.col(str(column_name))
+        if config.get("normalize_whitespace"):
+            column = F.regexp_replace(column, r"\s+", " ")
+        if config.get("trim"):
+            column = F.trim(column)
+        if config.get("lower"):
+            column = F.lower(column)
+        if config.get("upper"):
+            column = F.upper(column)
+        if config.get("empty_as_null"):
+            column = F.when(column == "", F.lit(None)).otherwise(column)
+        df = df.withColumn(str(column_name), column)
+    return df
+
+
+def _validate_columns(df: Any, columns: dict[str, Any], context: str) -> None:
+    available = set(getattr(df, "columns", ()) or ())
+    missing = sorted(str(column) for column in columns if str(column) not in available)
+    if missing:
+        raise ValueError(f"{context} references missing columns: {missing}")
+
+
+def _validate_column_mapping(df: Any, mapping: dict[str, Any]) -> None:
+    _validate_columns(df, mapping, "column_mapping")
+    existing = set(getattr(df, "columns", ()) or ())
+    targets = [str(target) for target in mapping.values()]
+    duplicates = sorted({target for target in targets if targets.count(target) > 1})
+    if duplicates:
+        raise ValueError(f"column_mapping has duplicate targets: {duplicates}")
+    reserved_targets = sorted(set(targets) & CONTROL_COLUMNS)
+    if reserved_targets:
+        raise ValueError(f"column_mapping cannot produce reserved control columns: {reserved_targets}")
+    collisions = sorted(
+        target
+        for source, target in ((str(source), str(target)) for source, target in mapping.items())
+        if target in existing and target != source
+    )
+    if collisions:
+        raise ValueError(f"column_mapping would collide with existing columns: {collisions}")
+
+
+def _apply_watermark_filter(df: Any, watermark_column: str | None, watermark_value: str | None) -> Any:
+    if not watermark_column or not watermark_value:
+        return df
+    from contractforge_databricks.watermark import render_watermark_filter_predicate
+
+    columns = tuple(part for part in watermark_column.split("|") if part)
+    _validate_columns(df, {column: True for column in columns}, "watermark_columns")
+    return df.where(render_watermark_filter_predicate(columns=columns, watermark_value=watermark_value))
+
+
+def _dict(value: object) -> dict[str, Any]:
+    return dict(value) if isinstance(value, dict) else {}
+
+
+def _contract_metadata(contract: SemanticContract) -> dict[str, Any]:
+    return dict(contract.operations.metadata or {}) if contract.operations and contract.operations.metadata else {}
+
+
+def _string_list(value: object) -> list[str]:
+    if value is None:
+        return []
+    if isinstance(value, str):
+        return [item.strip() for item in value.split(",") if item.strip()]
+    return [str(item) for item in value or ()]

contractforge_databricks/preparation/pyspark_staging.py

@@ -0,0 +1,70 @@
+"""PySpark write-mode staging helpers with lazy imports."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from contractforge_core.preparation import SCD2StageSpec, SnapshotStageSpec, resolved_hash_exclude_columns
+from contractforge_core.preparation import scd2_stage_spec_from_contract, snapshot_stage_spec_from_contract
+from contractforge_core.semantic import SemanticContract
+from contractforge_databricks.preparation.hashing import HASH_DELIMITER, HASH_NULL_SENTINEL, ROW_HASH_COLUMN
+
+
+def with_row_hash(df: Any, columns: tuple[str, ...], *, exclude: tuple[str, ...] = ()) -> Any:
+    from pyspark.sql import functions as F
+
+    included = tuple(column for column in columns if column not in set(exclude))
+    if not included:
+        raise ValueError("row hash requires at least one included column")
+    payload = [F.coalesce(F.col(column).cast("string"), F.lit(HASH_NULL_SENTINEL)) for column in included]
+    return df.withColumn(ROW_HASH_COLUMN, F.sha2(F.concat_ws(HASH_DELIMITER, *payload), 256))
+
+
+def prepare_snapshot_stage(df: Any, spec: SnapshotStageSpec) -> Any:
+    from pyspark.sql import functions as F
+
+    source_columns = tuple(column for column in spec.source_columns if column not in {"is_active", "deleted_at", "row_hash"})
+    staged = with_row_hash(df, source_columns)
+    return staged.withColumn(spec.is_active_column, F.lit(True)).withColumn(
+        spec.deleted_at_column,
+        F.lit(None).cast("timestamp"),
+    )
+
+
+def prepare_scd2_stage(df: Any, spec: SCD2StageSpec) -> Any:
+    from pyspark.sql import functions as F
+
+    staged = with_row_hash(df, spec.change_columns)
+    if spec.effective_from_column:
+        staged = staged.withColumn("valid_from", F.col(spec.effective_from_column).cast("timestamp"))
+    else:
+        staged = staged.withColumn("valid_from", F.current_timestamp())
+    staged = staged.withColumn("valid_to", F.lit(None).cast("timestamp"))
+    staged = staged.withColumn("is_current", F.lit(True))
+    staged = staged.withColumn("changed_columns", F.lit(None).cast("string"))
+    for key in spec.merge_keys:
+        staged = staged.withColumn(f"__merge_key_{key}", F.lit(None))
+    return staged
+
+
+def prepare_hash_diff_stage(df: Any, contract: SemanticContract) -> Any:
+    if contract.write.mode != "scd1_hash_diff":
+        raise ValueError("Hash-diff staging requires mode=scd1_hash_diff")
+    source_columns = tuple(str(column) for column in getattr(df, "columns", ()) or ())
+    hash_columns = source_columns if contract.write.hash_strategy == "all_columns_except" else contract.write.hash_keys
+    return with_row_hash(
+        df,
+        hash_columns,
+        exclude=resolved_hash_exclude_columns(contract),
+    )
+
+
+def apply_write_staging(df: Any, contract: SemanticContract) -> Any:
+    source_columns = tuple(str(column) for column in getattr(df, "columns", ()) or ())
+    if contract.write.mode == "scd1_hash_diff":
+        return prepare_hash_diff_stage(df, contract)
+    if contract.write.mode == "scd2_historical":
+        return prepare_scd2_stage(df, scd2_stage_spec_from_contract(contract, source_columns=source_columns))
+    if contract.write.mode == "snapshot_soft_delete":
+        return prepare_snapshot_stage(df, snapshot_stage_spec_from_contract(contract, source_columns=source_columns))
+    return df