overture-schema-codegen 0.1.1.dev0__py3-none-any.whl

overture/schema/codegen/pyspark/schema_builder.py

@@ -0,0 +1,187 @@
+"""Build StructType schema source from ModelSpec field trees."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..extraction.field import (
+    AnyScalar,
+    ArrayOf,
+    FieldShape,
+    LiteralScalar,
+    MapOf,
+    ModelRef,
+    NewTypeShape,
+    Primitive,
+    Scalar,
+    UnionRef,
+)
+from ..extraction.field_walk import enum_source
+from ..extraction.specs import FieldSpec, ModelSpec, UnionSpec
+from ..extraction.type_registry import get_type_mapping
+
+__all__ = [
+    "SHARED_TYPE_REFS",
+    "SchemaField",
+    "build_schema",
+]
+
+# Types whose base_type name maps to a _schema_structs.py StructType constant.
+# Reserved for types the codegen cannot walk (BBox is a plain class, not a
+# Pydantic BaseModel).  Pydantic BaseModels are inlined.
+SHARED_TYPE_REFS: dict[str, str] = {
+    "BBox": "BBOX_STRUCT",
+}
+
+# Literal and Enum fields both serialize as strings in Parquet.
+_STRING_FALLBACK = "StringType()"
+
+
+@dataclass(frozen=True, slots=True)
+class SchemaField:
+    """One field in the generated StructType.
+
+    Parameters
+    ----------
+    name
+        Column name.
+    type_expr
+        Spark type expression string (e.g. `"StringType()"`) or
+        a `_schema_structs.py` constant name.
+    """
+
+    name: str
+    type_expr: str
+
+
+def _spark_for_base(base_type: str, source_type: type | None) -> str:
+    """Return a Spark type expression for a primitive base type.
+
+    Tries `base_type` first, then falls back to `source_type.__name__`.
+    Returns `StringType()` when neither maps to a known Spark type.
+    """
+    mapping = get_type_mapping(base_type)
+    if mapping is not None and mapping.spark is not None:
+        return mapping.spark
+    if source_type is not None:
+        fallback = get_type_mapping(source_type.__name__)
+        if fallback is not None and fallback.spark is not None:
+            return fallback.spark
+    return _STRING_FALLBACK
+
+
+def _spark_for_scalar(scalar: Scalar) -> str:
+    """Map a `Scalar` variant to a Spark type expression.
+
+    `LiteralScalar` and `AnyScalar` serialize as strings. `Primitive`
+    scalars look up the type registry; enum primitives and BBox short-
+    circuit to strings / shared constants before the registry.
+    """
+    if isinstance(scalar, (LiteralScalar, AnyScalar)):
+        return _STRING_FALLBACK
+    if scalar.base_type in SHARED_TYPE_REFS:
+        return SHARED_TYPE_REFS[scalar.base_type]
+    if enum_source(scalar) is not None:
+        return _STRING_FALLBACK
+    return _spark_for_base(scalar.base_type, scalar.source_type)
+
+
+def _deduplicate_by_name(fields: list[FieldSpec]) -> list[FieldSpec]:
+    """Keep one FieldSpec per name, requiring every arm to agree on Spark type.
+
+    Union annotated_fields may contain the same field name declared by
+    multiple arms with different `FieldSpec`s -- a `Literal` discriminator
+    whose value differs per arm, or a field whose per-arm constraints
+    diverge (see `union_extraction.extract_union`). A columnar sink stores
+    one type per column name, so the schema needs exactly one entry. Two
+    same-named fields are compatible when they resolve to the SAME Spark
+    type -- the first-seen `FieldSpec`'s shape is kept (arbitrarily; the
+    column type is identical either way). Two same-named fields that resolve
+    to DIFFERENT Spark types cannot share one generated column, so this
+    always raises, whether the mismatch is numeric (a narrower int type vs a
+    float) or not.
+
+    Widening the two to their common type would often work in practice --
+    Spark and Parquet can promote a narrower numeric column to a wider one
+    (reading an int where the schema declares a double, say). It is forbidden
+    anyway: a widened column makes the union's type an implicit property
+    inferred from whichever arms happen to disagree, rather than a decision
+    stated in the model. Raising forces that decision to the surface at
+    generation instead of leaving it as a silent compatibility trap.
+    """
+    seen: dict[str, FieldSpec] = {}
+    for f in fields:
+        existing = seen.get(f.name)
+        if existing is None:
+            seen[f.name] = f
+            continue
+        spark_f, spark_existing = (
+            _shape_to_spark(f.shape),
+            _shape_to_spark(existing.shape),
+        )
+        if spark_f != spark_existing:
+            raise ValueError(
+                f"Union field {f.name!r} resolves to incompatible Spark "
+                f"types across arms ({spark_existing} vs {spark_f}); a "
+                "single Parquet column cannot represent both."
+            )
+    return list(seen.values())
+
+
+def _struct_type_expr(fields: list[FieldSpec]) -> str:
+    """Build an inline `StructType([...])` expression from a list of fields."""
+    if not fields:
+        raise ValueError(
+            "Cannot build a StructType for a model with no fields; an empty "
+            "struct column cannot carry data and signals an upstream "
+            "extraction problem."
+        )
+    parts = [
+        f'StructField("{f.name}", {_shape_to_spark(f.shape)}, True)' for f in fields
+    ]
+    return f"StructType([{', '.join(parts)}])"
+
+
+def _shape_to_spark(shape: FieldShape) -> str:
+    """Convert a FieldShape to a Spark type expression string."""
+    match shape:
+        case ArrayOf(element=element):
+            return f"ArrayType({_shape_to_spark(element)}, True)"
+        case NewTypeShape(inner=inner):
+            return _shape_to_spark(inner)
+        case ModelRef(model=m):
+            return _struct_type_expr(m.fields)
+        case UnionRef(union=u):
+            return _struct_type_expr(_deduplicate_by_name(u.fields))
+        case MapOf(key=k, value=v):
+            return f"MapType({_shape_to_spark(k)}, {_shape_to_spark(v)}, True)"
+        case Primitive() | LiteralScalar() | AnyScalar() as s:
+            return _spark_for_scalar(s)
+    raise TypeError(f"Unhandled FieldShape: {shape!r}")
+
+
+def build_schema(spec: ModelSpec) -> list[SchemaField]:
+    """Build schema fields for a feature spec.
+
+    Walks the field tree and maps types to Spark type expressions.
+    Recognizes shared types and emits fields in model order.
+
+    Parameters
+    ----------
+    spec
+        The feature spec to build schema fields for.
+
+    Returns
+    -------
+    list[SchemaField]
+        One entry per schema column in model order.
+    """
+    source_fields = (
+        _deduplicate_by_name(spec.fields)
+        if isinstance(spec, UnionSpec)
+        else spec.fields
+    )
+    return [
+        SchemaField(name=f.name, type_expr=_shape_to_spark(f.shape))
+        for f in source_fields
+    ]

overture/schema/codegen/pyspark/templates/_check_function.py.jinja2

@@ -0,0 +1,10 @@
+{%- macro check_function(c) -%}
+def {{ c.func_name }}() -> Check:
+    return Check(
+        field={{ c.field | py_literal }},
+        name={{ c.check_name | py_literal }},
+        expr={{ c.expr }},
+        shape=CheckShape.{{ c.shape }},
+        read_columns={{ c.read_columns | py_literal }},
+    )
+{% endmacro %}

overture/schema/codegen/pyspark/templates/model_module.py.jinja2

@@ -0,0 +1,83 @@
+{% from '_check_function.py.jinja2' import check_function -%}
+# This file is auto-generated by overture-schema-codegen. Do not edit.
+
+"""{{ model_title }} validation expression builders."""
+
+from __future__ import annotations
+
+from pyspark.sql import functions as F
+{% if spark_types %}
+from pyspark.sql.types import (
+{% for t in spark_types %}
+    {{ t }},
+{% endfor %}
+)
+{% endif %}
+{% if geometry_type %}
+from overture.schema.system.geometric import GeometryType
+
+{% endif %}
+from overture.schema.pyspark.check import Check, CheckShape, ModelValidation
+{% if schema_struct_refs %}
+from overture.schema.pyspark.expressions._schema_structs import (
+{% for r in schema_struct_refs %}
+    {{ r }},
+{% endfor %}
+)
+{% endif %}
+{% if column_pattern_fns %}
+from overture.schema.pyspark.expressions.column_patterns import (
+{% for f in column_pattern_fns %}
+    {{ f }},
+{% endfor %}
+)
+{% endif %}
+{% if constraint_expr_fns %}
+from overture.schema.pyspark.expressions.constraint_expressions import (
+{% for f in constraint_expr_fns %}
+    {{ f }},
+{% endfor %}
+)
+{% endif %}
+
+
+{% for c in check_functions %}
+{{ check_function(c) }}
+{% endfor %}
+
+def {{ model_name }}_checks() -> list[Check]:
+    """All validation checks for {{ model_name }}."""
+{% if check_functions %}
+    return [
+{% for c in check_functions %}
+        {{ c.func_name }}(),
+{% endfor %}
+    ]
+{% else %}
+    return []
+{% endif %}
+
+
+{{ schema_const_name }} = StructType(
+    [
+{%- for sf in schema_fields %}
+        StructField("{{ sf.name }}", {{ sf.type_expr }}, True),
+{%- endfor %}
+    ]
+)
+{% if geometry_types_literal %}
+
+GEOMETRY_TYPES: tuple[GeometryType, ...] = {{ geometry_types_literal }}
+{% endif %}
+
+ENTRY_POINT = "{{ entry_point }}"
+
+PARTITIONS: dict[str, str] = {{ partitions | py_literal }}
+
+MODEL_VALIDATION = ModelValidation(
+    schema={{ schema_const_name }},
+    checks={{ model_name }}_checks,
+{%- if geometry_types_literal %}
+    geometry_types=GEOMETRY_TYPES,
+{%- endif %}
+)

overture/schema/codegen/pyspark/templates/test_module.py.jinja2

@@ -0,0 +1,129 @@
+# Auto-generated — do not edit.
+
+"""Generated conformance tests for {{ model_name }}."""
+
+from __future__ import annotations
+
+import pytest
+from {{ expression_import }} import (
+    {{ model_name | upper }}_SCHEMA,
+    {{ model_name }}_checks,
+)
+from pyspark.sql import SparkSession
+
+from _support.harness import (
+    ValidationResults,
+    run_validation_pipeline,
+)
+{% if mutation_imports %}
+from _support.mutations import {{ mutation_imports | join(", ") }}
+{% endif %}
+{% if needs_set_at_path %}
+from _support.helpers import set_at_path
+{% endif %}
+from _support.scenarios import Scenario
+
+BASE_ROW_SPARSE: dict = {{ base_row_sparse }}
+
+
+BASE_ROW_POPULATED: dict = {{ base_row_populated }}
+
+
+SCENARIOS: list[Scenario] = [
+{% for entry in scenarios %}
+    Scenario(
+        {% for k, v in entry %}
+        {{ k }}={{ v }},
+        {% endfor %}
+    ),
+{% endfor %}
+]
+
+
+@pytest.fixture(scope="module")
+def checks() -> list:
+    return {{ model_name }}_checks()
+
+
+@pytest.fixture(scope="module")
+def sparse_results(spark: SparkSession, checks: list) -> ValidationResults:
+    return run_validation_pipeline(
+        spark,
+        {{ schema_name }},
+        checks,
+        BASE_ROW_SPARSE,
+        SCENARIOS,
+        model_name="{{ model_name }}",
+    )
+
+
+@pytest.fixture(scope="module")
+def populated_results(spark: SparkSession, checks: list) -> ValidationResults:
+    return run_validation_pipeline(
+        spark,
+        {{ schema_name }},
+        checks,
+        BASE_ROW_POPULATED,
+        SCENARIOS,
+        model_name="{{ model_name }}",
+    )
+
+
+def test_baseline_sparse(sparse_results: ValidationResults) -> None:
+    """Sparse base row passes every check the codegen produced.
+
+    Catches drift between base_row synthesis, schema_builder, and
+    check_builder -- if any of those produce output inconsistent with
+    the others (e.g. a check that rejects values the synthesizer emits
+    for required-only fields), the baseline fails here before any
+    scenario runs.
+    """
+    baseline = sparse_results.violations.get("{{ model_name }}::baseline", set())
+    assert baseline == set(), f"Sparse baseline has violations: {baseline}"
+
+
+def test_baseline_populated(populated_results: ValidationResults) -> None:
+    """Fully-populated base row passes every check the codegen produced.
+
+    Mirrors `test_baseline_sparse` but with all optional fields
+    filled, exercising codegen paths that only fire when a value is
+    present.
+    """
+    baseline = populated_results.violations.get("{{ model_name }}::baseline", set())
+    assert baseline == set(), f"Populated baseline has violations: {baseline}"
+
+
+@pytest.mark.parametrize("scenario", SCENARIOS, ids=lambda s: s.id)
+def test_scenario_sparse(
+    scenario: Scenario,
+    sparse_results: ValidationResults,
+) -> None:
+    _assert_scenario(scenario, sparse_results)
+
+
+@pytest.mark.parametrize("scenario", SCENARIOS, ids=lambda s: s.id)
+def test_scenario_populated(
+    scenario: Scenario,
+    populated_results: ValidationResults,
+) -> None:
+    _assert_scenario(scenario, populated_results)
+
+
+def _assert_scenario(
+    scenario: Scenario,
+    validation_results: ValidationResults,
+) -> None:
+    expected = (scenario.expected_field, scenario.expected_check)
+    if scenario.id in validation_results.skipped:
+        # An unbuildable scenario exercises nothing; fail loud rather than skip
+        # (a skip reads as a pass and hides codegen/scaffold gaps).
+        pytest.fail(
+            f"unbuildable scenario {scenario.id!r}: "
+            f"{validation_results.skipped[scenario.id]}"
+        )
+    valid_violations = validation_results.violations.get(f"{scenario.id}::valid", set())
+    assert expected not in valid_violations
+    invalid_violations = validation_results.violations.get(
+        f"{scenario.id}::invalid", set()
+    )
+    assert expected in invalid_violations

overture/schema/codegen/pyspark/test_data/__init__.py

@@ -0,0 +1,9 @@
+"""Test-data generation for the rendered PySpark conformance tests.
+
+Three modules cover three flavors of data:
+
+- `invalid_value`: constraint-violating values for triggering each check.
+- `base_row`: minimal and fully populated valid rows.
+- `scaffold`: sparse path scaffolds that supply the nested intermediates
+  (optional structs, arrays) a check's field path requires.
+"""