feat(pyspark): introduce package with codegen pipeline and CLI

Introduce overture-schema-pyspark, a runtime PySpark validation
package whose per-feature expression modules and conformance tests
are generated from the same Pydantic models that define the schema,
along with an `overture-validate` CLI.

Runtime (overture-schema-pyspark/src/overture/schema/pyspark/):

- check.py — Check, CheckShape, FeatureValidation dataclasses.
- schema_check.py — write-first comparison of Spark schemas against
  an expected StructType, with structural type matching and
  SchemaMismatch reporting.
- validate.py — public API: validate_feature(), evaluate_checks(),
  explain_errors(). The explain stage UNPIVOTs per-row check results
  into one row per violation, preserving all input columns for
  downstream join-back.
- cli.py — `overture-validate <parquet-or-directory>` runs the
  validation pipeline against a path of GeoParquet files. Output is
  one row per violation: feature ID, theme/type, failing field,
  check name, offending value. Single-pass evaluation keeps memory
  bounded for arbitrarily large inputs.
- expressions/ — shared runtime utilities (constraint_expressions,
  column_patterns, _schema_structs). Per-feature expression modules
  live under expressions/overture/ and are added by the codegen in
  a follow-up commit.
- tests/_support/ — conformance test infrastructure (scenarios,
  harness, helpers, mutations). The harness builds one DataFrame
  per feature, applies all scenarios as deterministic-UUID-tagged
  rows, runs validation once, and indexes violations back to
  scenario IDs — O(checks) rather than O(checks * scenarios).

CLI filtering options:

  --theme <theme>           limit to one theme
  --feature <feature>       limit to one feature type
  --skip-schema-check       run only constraint checks (no schema
                            comparison)
  --count-only              print violation counts per check rather
                            than rows
  --suppress <key>          suppress specific (feature, field, check)
                            triples per a YAML config

Codegen pipeline (overture-schema-codegen/src/.../pyspark/):

    FeatureSpec
        |
    constraint_dispatch.py   map constraints to descriptors
        |
    check_builder.py         walk FieldSpec -> CheckNode IR;
                             resolve array nesting, variant gating
        |
    schema_builder.py        FieldSpec -> SchemaField list
                             (StructType source)
        |
    renderer.py              CheckNode -> per-feature expression
                             module
    test_renderer.py         CheckNode -> per-feature conformance
                             test module
    synthetic.py             FeatureSpec -> BASE_ROW + invalid values
        |
    pipeline.py              orchestrate, return GeneratedModule list

The dispatch tables map every supported constraint (Ge/Gt/Le/Lt/
Interval, MinLen/MaxLen, StrippedConstraint, PatternConstraint,
UniqueItemsConstraint, GeometryTypeConstraint, JsonPointerConstraint,
RequireAnyOfConstraint, RadioGroupConstraint, RequireIfConstraint,
ForbidIfConstraint, MinFieldsSetConstraint), NewType (Country-
CodeAlpha2, LinearlyReferencedRange, RegionCode), and base type
(HttpUrl, EmailStr) to constraint_expressions check functions.

Discriminated unions (segment is the canonical hard case) split
into per-arm test files. The codegen handles arm splitting via
generate_arm_rows in synthetic.py and _filter_field_nodes_for_arm
in test_renderer.py.

The Makefile gains a `generate-pyspark` target and gates `check`
on it so a stale generation surfaces immediately. The CLI is exposed
as a `[project.scripts]` entry point so `overture-validate`
becomes available after `pip install` / `uv sync`.

Signed-off-by: Seth Fitzsimmons <seth@mojodna.net>

Author: sethfitz

Makefile

@@ -1,4 +1,4 @@
-.PHONY: default uv-sync check test-all test test-only docformat doctest doctest-only mypy mypy-only lint-only update-baselines
+.PHONY: default uv-sync clean-pyspark generate-pyspark check test-all test test-only docformat doctest doctest-only mypy mypy-only lint-only update-baselines
 
 TESTMON ?= --testmon
 
@@ -7,9 +7,22 @@ default: test-all
 install: uv-sync
 
 uv-sync:
-	@uv sync --all-packages 2> /dev/null
+	@uv sync --all-packages --all-extras 2> /dev/null
 
-check: uv-sync
+PYSPARK_EXPRESSIONS := packages/overture-schema-pyspark/src/overture/schema/pyspark/expressions/generated
+PYSPARK_GENERATED_TESTS := packages/overture-schema-pyspark/tests/generated
+
+clean-pyspark:
+	@rm -rf $(PYSPARK_EXPRESSIONS) $(PYSPARK_GENERATED_TESTS)
+
+generate-pyspark: uv-sync clean-pyspark
+	@uv run overture-codegen generate --format pyspark \
+		--output-dir $(PYSPARK_EXPRESSIONS) \
+		--test-output-dir $(PYSPARK_GENERATED_TESTS)
+	@uv run ruff check --fix --quiet $(PYSPARK_EXPRESSIONS) $(PYSPARK_GENERATED_TESTS)
+	@uv run ruff format --quiet $(PYSPARK_EXPRESSIONS) $(PYSPARK_GENERATED_TESTS)
+
+check: uv-sync generate-pyspark
 	@$(MAKE) -j test-only doctest-only lint-only mypy-only
 
 test-all: uv-sync

packages/overture-schema-codegen/docs/design.md

@@ -20,9 +20,25 @@ name = "overture-schema-codegen"
 overture-codegen = "overture.schema.codegen.cli:main"
 
 [tool.uv.sources]
+overture-schema-addresses-theme = { workspace = true }
+overture-schema-base-theme = { workspace = true }
+overture-schema-buildings-theme = { workspace = true }
 overture-schema-cli = { workspace = true }
 overture-schema-common = { workspace = true }
+overture-schema-divisions-theme = { workspace = true }
+overture-schema-places-theme = { workspace = true }
 overture-schema-system = { workspace = true }
+overture-schema-transportation-theme = { workspace = true }
+
+[dependency-groups]
+test = [
+    "overture-schema-addresses-theme",
+    "overture-schema-base-theme",
+    "overture-schema-buildings-theme",
+    "overture-schema-divisions-theme",
+    "overture-schema-places-theme",
+    "overture-schema-transportation-theme",
+]
 
 [tool.hatch.version]
 path = "src/overture/schema/codegen/__about__.py"

packages/overture-schema-codegen/docs/walkthrough.md

@@ -17,6 +17,7 @@
     FeatureSpec,
     is_model_class,
     is_union_alias,
+    partitions_from_tags,
 )
 from .extraction.union_extraction import extract_union
 from .layout.module_layout import (
@@ -26,12 +27,13 @@
     entry_point_module,
 )
 from .markdown.pipeline import generate_markdown_pages
+from .pyspark.pipeline import generate_pyspark_modules
 
 log = logging.getLogger(__name__)
 
 __all__ = ["cli"]
 
-_OUTPUT_FORMATS = ("markdown",)
+_OUTPUT_FORMATS = ("markdown", "pyspark")
 
 _FEATURE_FRONTMATTER = "---\nsidebar_position: 1\n---\n\n"
 
@@ -84,21 +86,29 @@ def list_models() -> None:
     "--output-dir",
     type=click.Path(path_type=Path),
     default=None,
-    help="Write output to directory (default: stdout)",
+    help="Write output files directly into this directory (default: stdout). "
+    "For pyspark, writes expression modules (*.py) and a _registry.py. "
+    "For markdown, writes theme subdirectories.",
+)
+@click.option(
+    "--test-output-dir",
+    type=click.Path(path_type=Path),
+    default=None,
+    help="Write test modules (test_*.py) into this directory (pyspark only).",
 )
 def generate(
     output_format: str,
     tags: tuple[str, ...],
     filters: tuple[str, ...],
     excludes: tuple[str, ...],
     output_dir: Path | None,
+    test_output_dir: Path | None,
 ) -> None:
     """Generate code/docs from discovered models."""
-    all_models = discover_models()
+    if output_format != "pyspark" and test_output_dir is not None:
+        raise click.UsageError("--test-output-dir is only valid with --format pyspark")
 
-    # Schema root from ALL entry points (before tag filters).
-    module_paths = [entry_point_module(k.entry_point) for k in all_models]
-    schema_root = compute_schema_root(module_paths)
+    all_models = discover_models()
 
     models = filter_models(all_models, build_selector(tags, filters, excludes))
 
@@ -107,18 +117,27 @@ def generate(
 
     feature_specs: list[FeatureSpec] = []
     for key, entry in models.items():
+        partitions = partitions_from_tags(key.tags)
         if is_model_class(entry):
-            feature_specs.append(extract_model(entry, entry_point=key.entry_point))
+            feature_specs.append(
+                extract_model(entry, entry_point=key.entry_point, partitions=partitions)
+            )
         elif is_union_alias(entry):
             feature_specs.append(
                 extract_union(
                     entry_point_class(key.entry_point),
                     entry,
                     entry_point=key.entry_point,
+                    partitions=partitions,
                 )
             )
 
-    _generate_markdown(feature_specs, schema_root, output_dir)
+    if output_format == "pyspark":
+        _generate_pyspark(feature_specs, output_dir, test_output_dir)
+    else:
+        module_paths = [entry_point_module(k.entry_point) for k in all_models]
+        schema_root = compute_schema_root(module_paths)
+        _generate_markdown(feature_specs, schema_root, output_dir)
 
 
 def _generate_markdown(
@@ -141,6 +160,24 @@ def _generate_markdown(
         _write_category_files(output_dir, all_paths, feature_paths)
 
 
+def _generate_pyspark(
+    feature_specs: list[FeatureSpec],
+    output_dir: Path | None,
+    test_output_dir: Path | None = None,
+) -> None:
+    """Generate PySpark validation modules.
+
+    Output is syntactically valid Python; we assume a code formatter runs
+    over the written directories afterwards to match existing conventions.
+    """
+    modules = generate_pyspark_modules(feature_specs)
+    for mod in modules.source:
+        _write_output(mod.content, output_dir, mod.path)
+    if test_output_dir is not None:
+        for mod in modules.test:
+            _write_output(mod.content, test_output_dir, mod.path)
+
+
 def _ancestor_dirs(paths: set[PurePosixPath]) -> set[PurePosixPath]:
     """Collect all ancestor directories for a set of file paths."""
     dirs: set[PurePosixPath] = set()

packages/overture-schema-codegen/pyproject.toml

@@ -0,0 +1,172 @@
+"""Tree-shaped IR for model field types.
+
+`FieldShape` is a discriminated union -- `Primitive`, `LiteralScalar`,
+`AnyScalar`, `ModelRef`, `UnionRef`, `ArrayOf`, `MapOf`, `NewTypeShape`
+-- nested to describe arbitrary list / dict / NewType wrapping. Each
+variant carries its own constraints (where meaningful), and walkers
+encounter each constraint at the layer it targets.
+
+The three terminal scalar variants (`Primitive`, `LiteralScalar`,
+`AnyScalar`) are grouped under the `Scalar` type alias for consumers
+that only need to ask "is this a leaf?".
+
+`NewTypeShape` wraps an inner shape, so its position relative to
+`ArrayOf` is structural: `NewTypeShape(inner=ArrayOf(...))` is a
+NewType over `list[X]`, while `ArrayOf(element=NewTypeShape(...))`
+is a list of NewType-wrapped values. Consumers pattern-match on
+shape to distinguish the two.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeAlias
+
+if TYPE_CHECKING:
+    from .specs import ModelSpec, UnionSpec
+
+__all__ = [
+    "AnyScalar",
+    "ArrayOf",
+    "ConstraintSource",
+    "FieldShape",
+    "LiteralScalar",
+    "MapOf",
+    "ModelRef",
+    "NewTypeShape",
+    "Primitive",
+    "Scalar",
+    "UnionRef",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class ConstraintSource:
+    """A constraint paired with the NewType that contributed it.
+
+    `source_ref` and `source_name` identify the NewType that declared
+    the constraint; both are `None` for constraints contributed directly
+    on a field annotation rather than through a NewType. `constraint`
+    is the raw metadata object from `Annotated[..., constraint]`.
+    """
+
+    source_ref: object | None
+    source_name: str | None
+    constraint: object
+
+
+@dataclass(frozen=True, slots=True)
+class Primitive:
+    """Terminal type with a registry lookup key.
+
+    Covers primitives (`int32`, `str`), enums, Pydantic built-ins
+    (`HttpUrl`, `EmailStr`), and `BaseModel` subclasses that weren't
+    resolved to a `ModelRef` (e.g. when no `model_resolver` was
+    supplied).
+    """
+
+    base_type: str
+    source_type: type | None = None
+    constraints: tuple[ConstraintSource, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class LiteralScalar:
+    """`Literal[X, ...]` terminal."""
+
+    values: tuple[object, ...]
+    constraints: tuple[ConstraintSource, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class AnyScalar:
+    """`typing.Any` terminal."""
+
+    constraints: tuple[ConstraintSource, ...] = ()
+
+
+Scalar: TypeAlias = Primitive | LiteralScalar | AnyScalar
+"""Terminal shape: a value that doesn't wrap another shape.
+
+Consumers that just need "is this a leaf?" check `isinstance(x, Scalar)`;
+consumers that need terminal-specific data narrow to a variant.
+"""
+
+
+@dataclass(frozen=True, slots=True)
+class ModelRef:
+    """Reference to a Pydantic sub-model.
+
+    `starts_cycle` marks the back-edge of a cycle in the model graph;
+    consumers that recurse into models must stop at cycle starts.
+    """
+
+    model: ModelSpec
+    starts_cycle: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class UnionRef:
+    """Reference to a discriminated union of models."""
+
+    union: UnionSpec
+
+
+@dataclass(frozen=True, slots=True)
+class ArrayOf:
+    """Sequence of values sharing a single element shape.
+
+    Nested arrays are nested `ArrayOf` instances; there is no numeric
+    depth field. `constraints` carries array-level validation rules
+    (length, uniqueness). Per-element constraints live on `element`
+    and its descendants.
+    """
+
+    element: FieldShape
+    constraints: tuple[ConstraintSource, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class MapOf:
+    """Mapping from a key shape to a value shape.
+
+    `constraints` carries map-level validation rules. Per-key and
+    per-value constraints live on `key` / `value` respectively.
+    """
+
+    key: FieldShape
+    value: FieldShape
+    constraints: tuple[ConstraintSource, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class NewTypeShape:
+    """A NewType wrapper around an inner shape.
+
+    Position relative to other wrappers is meaningful:
+    `NewTypeShape(inner=ArrayOf(...))` is a NewType over `list[X]`;
+    `ArrayOf(element=NewTypeShape(...))` is a list of NewType-wrapped
+    values. Consumers distinguish the two by pattern, not a numeric
+    offset.
+
+    Constraints contributed by the NewType chain attach to the
+    `Scalar` / `ArrayOf` / `MapOf` layer they target, not to the
+    wrapper itself. `name` and `ref` identify the NewType for linking
+    without owning constraint state.
+    """
+
+    name: str
+    ref: object
+    inner: FieldShape
+
+
+FieldShape: TypeAlias = (
+    Primitive
+    | LiteralScalar
+    | AnyScalar
+    | ModelRef
+    | UnionRef
+    | ArrayOf
+    | MapOf
+    | NewTypeShape
+)