spark-joinery 0.0.1__tar.gz

spark_joinery-0.0.1/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.3
+Name: spark-joinery
+Version: 0.0.1
+Summary: Schema safety for pyspark transformations
+Author: Connor Charles
+Author-email: Connor Charles <ccharles.gb@gmail.com>
+Requires-Dist: pyspark>=4.2.0
+Requires-Dist: typer>=0.19.2 ; extra == 'cli'
+Requires-Dist: pydantic>=2.0.0 ; extra == 'pydantic'
+Requires-Python: >=3.13
+Provides-Extra: cli
+Provides-Extra: pydantic

spark_joinery-0.0.1/pyproject.toml

@@ -0,0 +1,41 @@
+[project]
+name = "spark-joinery"
+version = "0.0.1"
+description = "Schema safety for pyspark transformations"
+requires-python = ">=3.13"
+dependencies = ["pyspark>=4.2.0"]
+
+[[project.authors]]
+name = "Connor Charles"
+email = "ccharles.gb@gmail.com"
+
+[project.optional-dependencies]
+cli = ["typer>=0.19.2"]
+pydantic = ["pydantic>=2.0.0"]
+
+[project.scripts]
+spark-joinery = "spark_joinery:main"
+
+[tool.pyrefly]
+project-includes = [
+    "**/*.py*",
+    "**/*.ipynb",
+]
+search-path = [
+    "src",
+    "examples",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.27,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "ipykernel>=7.3.0",
+    "pydantic>=2.0.0",
+    "pyrefly>=1.2.0",
+    "pytest>=9.1.0",
+    "ruff>=0.15.17",
+    "zensical>=0.0.45",
+]

spark_joinery-0.0.1/pyproject.toml.orig

@@ -0,0 +1,46 @@
+[project]
+name = "spark-joinery"
+version = "0.0.1"
+description = "Schema safety for pyspark transformations"
+authors = [
+    { name = "Connor Charles", email = "ccharles.gb@gmail.com" }
+]
+requires-python = ">=3.13"
+dependencies = [
+    "pyspark>=4.2.0",
+]
+
+[tool.pyrefly]
+project-includes = [
+    "**/*.py*",
+    "**/*.ipynb",
+]
+search-path = [
+    "src",
+    "examples",
+]
+
+[project.optional-dependencies]
+cli = [
+    "typer>=0.19.2",
+]
+pydantic = [
+    "pydantic>=2.0.0",
+]
+
+[project.scripts]
+spark-joinery = "spark_joinery:main"
+
+[build-system]
+requires = ["uv_build>=0.9.27,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "ipykernel>=7.3.0",
+    "pydantic>=2.0.0",
+    "pyrefly>=1.2.0",
+    "pytest>=9.1.0",
+    "ruff>=0.15.17",
+    "zensical>=0.0.45",
+]

spark_joinery-0.0.1/src/spark_joinery/__init__.py

@@ -0,0 +1,20 @@
+from .collection import Collection
+from .dependencies import Context, PipelineContext
+from .pipeline import (
+    ExecutablePipeline,
+    Pipeline,
+    PipelineExecutionError,
+    Step,
+)
+from .transform import transform
+
+__all__ = [
+    "Collection",
+    "Context",
+    "ExecutablePipeline",
+    "Pipeline",
+    "PipelineExecutionError",
+    "PipelineContext",
+    "Step",
+    "transform",
+]

spark_joinery-0.0.1/src/spark_joinery/collection.py

@@ -0,0 +1,53 @@
+from typing import Any, Callable, ParamSpec, TypeVar, overload
+
+from .schemas import CoercionMode
+from .transform import TransformSpec, _inspect_transform, _wrap_transform
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class Collection:
+    def __init__(self) -> None:
+        self._specs: dict[Callable[..., Any], TransformSpec] = {}
+
+    @overload
+    def transform(
+        self,
+        f: Callable[P, R],
+        *,
+        validate_input: CoercionMode | None = "project_all",
+        validate_output: CoercionMode | None = "project_all",
+    ) -> Callable[P, R]: ...
+
+    @overload
+    def transform(
+        self,
+        f: None = None,
+        *,
+        validate_input: CoercionMode | None = "project_all",
+        validate_output: CoercionMode | None = "project_all",
+    ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+    def transform(
+        self,
+        f: Callable[P, R] | None = None,
+        *,
+        validate_input: CoercionMode | None = "project_all",
+        validate_output: CoercionMode | None = "project_all",
+    ):
+        def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+            spec = _inspect_transform(fn)
+            wrapper = _wrap_transform(
+                fn,
+                spec,
+                validate_input=validate_input,
+                validate_output=validate_output,
+            )
+            self._specs[wrapper] = spec
+            return wrapper
+
+        if f is None:
+            return decorator
+
+        return decorator(f)

spark_joinery-0.0.1/src/spark_joinery/dependencies.py

@@ -0,0 +1,28 @@
+from typing import Any, Sequence
+
+
+class Context:
+    """Marks a transform parameter as resolved from a PipelineContext."""
+
+    def __init__(self, type_: type | None = None) -> None:
+        self.type = type_
+
+
+class PipelineContext:
+    """Registry of dependency values, looked up by type."""
+
+    def __init__(self, *, values: Sequence[Any] = ()) -> None:
+        self._by_type: dict[type, Any] = {}
+        for value in values:
+            value_type = type(value)
+            if value_type in self._by_type:
+                raise ValueError(
+                    f"duplicate dependency type {value_type!r} in PipelineContext"
+                )
+            self._by_type[value_type] = value
+
+    def resolve(self, param_type: type, marker: "Context") -> Any:
+        lookup_type = marker.type or param_type
+        if lookup_type not in self._by_type:
+            raise KeyError(f"no dependency registered for type {lookup_type!r}")
+        return self._by_type[lookup_type]

spark_joinery-0.0.1/src/spark_joinery/fixtures.py

@@ -0,0 +1,23 @@
+from collections.abc import Sequence
+from typing import TypeVar
+from pyspark.sql import DataFrame, SparkSession
+from spark_joinery import schemas
+
+T = TypeVar("T")
+
+
+def get_dataframe(
+    spark: SparkSession, row_type: type[T], rows: Sequence[T]
+) -> DataFrame:
+    for index, row in enumerate(rows):
+        if not isinstance(row, row_type):
+            raise ValueError(
+                f"Row {index} of type {row.__class__.__name__}. Expected type {row_type.__name__}"
+            )
+    schema = schemas.get_spark_schema_from_model(row_type)
+
+    serialized_rows = [
+        row.model_dump(mode="python") if hasattr(row, "model_dump") else row
+        for row in rows
+    ]
+    return spark.createDataFrame(serialized_rows, schema)

spark_joinery-0.0.1/src/spark_joinery/pipeline.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Sequence
+
+from pyspark.sql import DataFrame, SparkSession, types
+
+from .collection import Collection
+from .dependencies import Context, PipelineContext
+from .schemas import CoercionMode
+
+Transform = Callable[..., DataFrame | None]
+
+
+class PipelineExecutionError(RuntimeError):
+    pass
+
+
+@dataclass(eq=False)
+class Step:
+    name: str
+    transform: Transform
+    _pipeline: Pipeline
+    _input_schemas: dict[str, types.StructType]
+    _output_schema: types.StructType | None
+    _spark_parameter: str | None
+    _context_parameters: dict[str, tuple[type, Context]]
+    _upstream_steps: list[Step]
+    _explicit_bindings: dict[Step, str]
+
+
+@dataclass(frozen=True)
+class _ExecutionStep:
+    step: Step
+    dataframe_bindings: tuple[tuple[str, Step], ...]
+
+
+class ExecutablePipeline:
+    def __init__(self, execution_steps: Sequence[_ExecutionStep]):
+        self._execution_steps = tuple(execution_steps)
+
+    def run(
+        self, spark: SparkSession, context: PipelineContext | None = None
+    ) -> dict[str, DataFrame]:
+        if not isinstance(spark, SparkSession):
+            raise TypeError("run() requires a SparkSession")
+
+        outputs: dict[str, DataFrame] = {}
+        for execution_step in self._execution_steps:
+            step = execution_step.step
+            arguments: dict[str, Any] = {
+                parameter_name: outputs[upstream_step.name]
+                for parameter_name, upstream_step in execution_step.dataframe_bindings
+            }
+            if step._spark_parameter is not None:
+                arguments[step._spark_parameter] = spark
+
+            for parameter_name, (
+                param_type,
+                marker,
+            ) in step._context_parameters.items():
+                if context is None:
+                    raise PipelineExecutionError(
+                        f"step '{step.name}' requires a PipelineContext but none was provided"
+                    )
+                try:
+                    arguments[parameter_name] = context.resolve(param_type, marker)
+                except KeyError as error:
+                    raise PipelineExecutionError(
+                        f"step '{step.name}' failed to resolve dependency"
+                    ) from error
+
+            try:
+                result = step.transform(**arguments)
+                if step._output_schema is not None and not isinstance(
+                    result, DataFrame
+                ):
+                    raise TypeError(
+                        f"step '{step.name}' returned a non-DataFrame value"
+                    )
+            except Exception as error:
+                if isinstance(error, PipelineExecutionError):
+                    raise
+                raise PipelineExecutionError(
+                    f"Pipeline step '{step.name}' failed"
+                ) from error
+
+            if isinstance(result, DataFrame):
+                outputs[step.name] = result
+
+        return outputs
+
+
+class Pipeline:
+    def __init__(self, collections: Sequence[Collection] = ()):
+        self._collection = Collection()
+        self._collections: tuple[Collection, ...] = (self._collection, *collections)
+        self._steps: dict[str, Step] = {}
+        self._validated = False
+        self._executable: ExecutablePipeline | None = None
+
+    def transform(
+        self,
+        f=None,
+        *,
+        validate_input: CoercionMode | None = "project_all",
+        validate_output: CoercionMode | None = "project_all",
+    ):
+        return self._collection.transform(
+            f,
+            validate_input=validate_input,
+            validate_output=validate_output,
+        )
+
+    def _resolve_spec(self, transform: Transform):
+        for collection in self._collections:
+            spec = collection._specs.get(transform)
+            if spec is not None:
+                return spec
+        name = getattr(transform, "__name__", repr(transform))
+        raise TypeError(f"'{name}' is not registered in this pipeline's collections")
+
+    def add_step(self, transform: Transform, name: str) -> Step:
+        self._ensure_mutable()
+        if name in self._steps:
+            raise ValueError(f"step name '{name}' is already registered")
+
+        spec = self._resolve_spec(transform)
+        input_schemas = spec.input_schemas
+        output_schema = spec.output_schema
+
+        spark_parameter = spec.spark_parameter
+        parameter_names = set(inspect.signature(transform).parameters)
+        supported_parameters = (
+            set(input_schemas)
+            | ({spark_parameter} if spark_parameter is not None else set())
+            | set(spec.context_parameters)
+        )
+        unsupported_parameters = parameter_names - supported_parameters
+        if unsupported_parameters:
+            unsupported = sorted(unsupported_parameters)[0]
+            raise TypeError(
+                "pipeline steps only support SparkSession, annotated DataFrame, "
+                f"and Context-annotated parameters; unsupported parameter '{unsupported}'"
+            )
+
+        if not input_schemas and spark_parameter is None:
+            raise TypeError("source step must declare a SparkSession")
+
+        step = Step(
+            name=name,
+            transform=transform,
+            _pipeline=self,
+            _input_schemas=input_schemas,
+            _output_schema=output_schema,
+            _spark_parameter=spark_parameter,
+            _context_parameters=spec.context_parameters,
+            _upstream_steps=[],
+            _explicit_bindings={},
+        )
+        self._steps[name] = step
+        return step
+
+    def connect(
+        self, upstream: Step, downstream: Step, *, param: str | None = None
+    ) -> None:
+        self._ensure_mutable()
+        self._ensure_owned(upstream)
+        self._ensure_owned(downstream)
+        if upstream not in downstream._upstream_steps:
+            downstream._upstream_steps.append(upstream)
+        if param is not None:
+            for other_upstream, other_param in downstream._explicit_bindings.items():
+                if other_param == param and other_upstream is not upstream:
+                    raise ValueError(
+                        f"step '{downstream.name}' parameter '{param}' is already "
+                        f"bound to upstream step '{other_upstream.name}'"
+                    )
+            downstream._explicit_bindings[upstream] = param
+
+    def connect_many(
+        self,
+        upstream_steps: Sequence[Step],
+        downstream: Step,
+    ) -> None:
+        self._ensure_mutable()
+        if not upstream_steps:
+            raise ValueError("connect_many requires at least one upstream step")
+        for upstream in upstream_steps:
+            self.connect(upstream, downstream)
+
+    def validate(self) -> ExecutablePipeline:
+        if self._executable is not None:
+            return self._executable
+
+        execution_steps: list[_ExecutionStep] = []
+        for step in self._topological_order():
+            bindings: list[tuple[str, Step]] = []
+            matched_upstream: set[Step] = set()
+            for parameter_name, expected_schema in step._input_schemas.items():
+                explicit_upstream = next(
+                    (
+                        upstream
+                        for upstream, bound_param in step._explicit_bindings.items()
+                        if bound_param == parameter_name
+                    ),
+                    None,
+                )
+                if explicit_upstream is not None:
+                    if explicit_upstream._output_schema != expected_schema:
+                        raise ValueError(
+                            f"step '{step.name}' parameter '{parameter_name}' is "
+                            f"explicitly bound to '{explicit_upstream.name}' but its "
+                            "output schema does not match"
+                        )
+                    bindings.append((parameter_name, explicit_upstream))
+                    matched_upstream.add(explicit_upstream)
+                    continue
+
+                candidates = [
+                    upstream
+                    for upstream in step._upstream_steps
+                    if upstream not in step._explicit_bindings
+                ]
+                matches = [
+                    upstream
+                    for upstream in candidates
+                    if upstream._output_schema == expected_schema
+                ]
+                if not matches:
+                    raise ValueError(
+                        f"step '{step.name}' parameter '{parameter_name}' has no "
+                        "upstream step provides a matching schema"
+                    )
+                if len(matches) > 1:
+                    names = ", ".join(upstream.name for upstream in matches)
+                    raise ValueError(
+                        f"step '{step.name}' has ambiguous upstream schema for "
+                        f"parameter '{parameter_name}': {names}"
+                    )
+                upstream = matches[0]
+                bindings.append((parameter_name, upstream))
+                matched_upstream.add(upstream)
+
+            unmatched = [
+                upstream
+                for upstream in step._upstream_steps
+                if upstream not in matched_upstream
+            ]
+            if unmatched:
+                names = ", ".join(upstream.name for upstream in unmatched)
+                raise ValueError(
+                    f"step '{step.name}' has upstream output that does not match "
+                    f"a parameter: {names}"
+                )
+
+            execution_steps.append(_ExecutionStep(step, tuple(bindings)))
+
+        self._validated = True
+        self._executable = ExecutablePipeline(execution_steps)
+        return self._executable
+
+    def _ensure_mutable(self) -> None:
+        if self._validated:
+            raise RuntimeError("pipeline has already been validated")
+
+    def _ensure_owned(self, step: Step) -> None:
+        if step._pipeline is not self:
+            raise ValueError("steps must belong to the same pipeline")
+
+    def _topological_order(self) -> list[Step]:
+        states: dict[Step, int] = {}
+        ordered: list[Step] = []
+
+        def visit(step: Step) -> None:
+            state = states.get(step, 0)
+            if state == 1:
+                raise ValueError("pipeline contains a cycle")
+            if state == 2:
+                return
+            states[step] = 1
+            for upstream in step._upstream_steps:
+                visit(upstream)
+            states[step] = 2
+            ordered.append(step)
+
+        for step in self._steps.values():
+            visit(step)
+        return ordered

spark_joinery-0.0.1/src/spark_joinery/schemas.py

@@ -0,0 +1,188 @@
+from dataclasses import fields, is_dataclass
+from functools import partial
+from typing import Any, Callable, TypeVar, Literal
+
+from pyspark.errors import PySparkAssertionError
+from pyspark.sql import Column, DataFrame, functions as F, types
+from pyspark.testing import assertSchemaEqual
+
+from . import type_inspection
+
+T = TypeVar("T")
+
+CoercionMode = Literal["coerce", "project_all", "project", "strict", "strict_null"]
+
+
+def is_pydantic_model(klass: Any) -> bool:
+    # Avoid hard dependency on pydantic; detect by BaseModel class attributes.
+    return (
+        isinstance(klass, type)
+        and hasattr(klass, "model_fields")
+        and isinstance(getattr(klass, "model_fields"), dict)
+    )
+
+
+def is_schema_model(klass: Any) -> bool:
+    return is_dataclass(klass) or is_pydantic_model(klass)
+
+
+def _get_model_fields(klass: type[Any]) -> list[tuple[str, Any]]:
+    if is_dataclass(klass):
+        return [(field.name, field.type) for field in fields(klass)]
+
+    if is_pydantic_model(klass):
+        model_fields = getattr(klass, "model_fields")
+        return [
+            (name, field_info.annotation)
+            for name, field_info in model_fields.items()
+            if field_info.annotation is not None
+        ]
+
+    raise ValueError(f"{klass.__name__} is neither a dataclass nor a pydantic model")
+
+
+def _get_spark_field_type(field_type: Any) -> types.DataType:
+    normalized_type = type_inspection.normalize_python_type(field_type)
+
+    if is_schema_model(normalized_type):
+        return get_spark_schema_from_model(normalized_type)
+
+    if type_inspection.is_list(normalized_type):
+        element_type = type_inspection.get_list_element_type(normalized_type)
+        element_spark_type = _get_spark_field_type(element_type)
+        return types.ArrayType(element_spark_type, True)
+
+    return type_inspection.get_spark_type_from_python_type(normalized_type)
+
+
+def get_spark_schema_from_model(klass: type[T]) -> types.StructType:
+    if not is_schema_model(klass):
+        raise ValueError(
+            f"{klass.__name__} is neither a dataclass nor a pydantic model"
+        )
+
+    struct_fields = [
+        types.StructField(name, _get_spark_field_type(field_type), True)
+        for name, field_type in _get_model_fields(klass)
+    ]
+
+    return types.StructType(struct_fields)
+
+
+def get_spark_schema_from_dataclass(klass: type[T]) -> types.StructType:
+    if not is_dataclass(klass):
+        raise ValueError(f"{klass.__name__} is not a dataclass")
+
+    return get_spark_schema_from_model(klass)
+
+
+def dataframe_is_schema(dataframe: DataFrame, klass: type[T]) -> bool:
+    expected_schema = get_spark_schema_from_dataclass(klass)
+    return dataframe.schema == expected_schema
+
+
+def dataframe_is_model_schema(dataframe: DataFrame, klass: type[T]) -> bool:
+    expected_schema = get_spark_schema_from_model(klass)
+    return dataframe.schema == expected_schema
+
+
+def _coerce_strict(
+    dataframe: DataFrame, schema: types.StructType, *, ignore_nullable: bool
+) -> DataFrame:
+    try:
+        assertSchemaEqual(dataframe.schema, schema, ignoreNullable=ignore_nullable)
+    except PySparkAssertionError as e:
+        raise ValueError("Schema mismatch") from e
+    return dataframe
+
+
+def _build_projected_column(
+    column: Column,
+    source_type: types.DataType,
+    target_type: types.DataType,
+    *,
+    cast: bool,
+    recurse: bool,
+) -> Column:
+    if recurse and isinstance(target_type, types.StructType):
+        if not isinstance(source_type, types.StructType):
+            raise ValueError(f"Expected struct type but found {source_type}")
+
+        source_fields = {field.name: field.dataType for field in source_type.fields}
+        nested_columns = []
+        for target_field in target_type.fields:
+            if target_field.name not in source_fields:
+                raise ValueError(f"Missing field '{target_field.name}'")
+            nested_column = _build_projected_column(
+                column.getField(target_field.name),
+                source_fields[target_field.name],
+                target_field.dataType,
+                cast=cast,
+                recurse=recurse,
+            )
+            nested_columns.append(nested_column.alias(target_field.name))
+        return F.struct(*nested_columns)
+
+    if recurse and isinstance(target_type, types.ArrayType):
+        if not isinstance(source_type, types.ArrayType):
+            raise ValueError(f"Expected array type but found {source_type}")
+
+        return F.transform(
+            column,
+            lambda element: _build_projected_column(
+                element,
+                source_type.elementType,
+                target_type.elementType,
+                cast=cast,
+                recurse=recurse,
+            ),
+        )
+
+    if source_type == target_type:
+        return column
+
+    if cast:
+        return column.cast(target_type)
+
+    raise ValueError(f"Type mismatch: expected {target_type}, found {source_type}")
+
+
+def _project_fields(
+    dataframe: DataFrame, schema: types.StructType, *, cast: bool, recurse: bool
+) -> DataFrame:
+    source_fields = {field.name: field.dataType for field in dataframe.schema.fields}
+
+    missing_fields = [
+        field.name for field in schema.fields if field.name not in source_fields
+    ]
+    if missing_fields:
+        raise ValueError(f"Missing columns: {missing_fields}")
+
+    columns = [
+        _build_projected_column(
+            F.col(field.name),
+            source_fields[field.name],
+            field.dataType,
+            cast=cast,
+            recurse=recurse,
+        ).alias(field.name)
+        for field in schema.fields
+    ]
+    return dataframe.select(*columns)
+
+
+_MODE_HANDLERS: dict[
+    CoercionMode, Callable[[DataFrame, types.StructType], DataFrame]
+] = {
+    "strict": partial(_coerce_strict, ignore_nullable=True),
+    "strict_null": partial(_coerce_strict, ignore_nullable=False),
+    "project": partial(_project_fields, cast=False, recurse=False),
+    "project_all": partial(_project_fields, cast=False, recurse=True),
+    "coerce": partial(_project_fields, cast=True, recurse=True),
+}
+
+
+def coerce_dataframe(
+    dataframe: DataFrame, schema: types.StructType, mode: CoercionMode
+) -> DataFrame:
+    return _MODE_HANDLERS[mode](dataframe, schema)

spark_joinery-0.0.1/src/spark_joinery/transform.py

@@ -0,0 +1,212 @@
+from dataclasses import dataclass
+from functools import wraps
+import inspect
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    ParamSpec,
+    TypeVar,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+    overload,
+)
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import types
+
+from . import schemas
+from .dependencies import Context
+from .schemas import CoercionMode
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def _get_annotated_dataframe_schema_model(annotation: Any) -> type[Any] | None:
+    if get_origin(annotation) is not Annotated:
+        return None
+
+    annotated_args = get_args(annotation)
+    if len(annotated_args) < 2:
+        return None
+
+    base_type = annotated_args[0]
+    metadata = annotated_args[1:]
+    if base_type is not DataFrame:
+        return None
+
+    for metadata_value in metadata:
+        if isinstance(metadata_value, type) and schemas.is_schema_model(metadata_value):
+            return metadata_value
+
+    return None
+
+
+def _get_context_marker(annotation: Any) -> tuple[type, Context] | None:
+    if get_origin(annotation) is not Annotated:
+        return None
+
+    annotated_args = get_args(annotation)
+    if len(annotated_args) < 2:
+        return None
+
+    base_type = annotated_args[0]
+    metadata = annotated_args[1:]
+    for metadata_value in metadata:
+        if isinstance(metadata_value, Context):
+            return base_type, metadata_value
+
+    return None
+
+
+@dataclass(frozen=True)
+class TransformSpec:
+    input_schemas: dict[str, types.StructType]
+    output_schema: types.StructType | None
+    spark_parameter: str | None
+    context_parameters: dict[str, tuple[type, Context]]
+
+
+def _get_spark_parameter(f: Any) -> str | None:
+    type_hints = get_type_hints(f)
+    spark_parameters = [
+        parameter_name
+        for parameter_name in inspect.signature(f).parameters
+        if type_hints.get(parameter_name) is SparkSession
+    ]
+    if len(spark_parameters) > 1:
+        raise TypeError("transform may declare only one SparkSession parameter")
+    return spark_parameters[0] if spark_parameters else None
+
+
+def _inspect_transform(f: Any) -> TransformSpec:
+    signature = inspect.signature(f)
+    type_hints = get_type_hints(f, include_extras=True)
+    input_schemas: dict[str, types.StructType] = {}
+    context_parameters: dict[str, tuple[type, Context]] = {}
+
+    for parameter_name in signature.parameters:
+        parameter_type = type_hints.get(parameter_name)
+        if parameter_type is None:
+            continue
+
+        dataframe_schema_model = _get_annotated_dataframe_schema_model(parameter_type)
+        if dataframe_schema_model is not None:
+            input_schemas[parameter_name] = schemas.get_spark_schema_from_model(
+                dataframe_schema_model
+            )
+            continue
+
+        context_marker = _get_context_marker(parameter_type)
+        if context_marker is not None:
+            context_parameters[parameter_name] = context_marker
+
+    output_schema = None
+    return_type = type_hints.get("return")
+    if return_type is not None:
+        dataframe_schema_model = _get_annotated_dataframe_schema_model(return_type)
+        if dataframe_schema_model is not None:
+            output_schema = schemas.get_spark_schema_from_model(dataframe_schema_model)
+
+    return TransformSpec(
+        input_schemas=input_schemas,
+        output_schema=output_schema,
+        spark_parameter=_get_spark_parameter(f),
+        context_parameters=context_parameters,
+    )
+
+
+def _wrap_transform(
+    fn: Callable[P, R],
+    spec: TransformSpec,
+    *,
+    validate_input: CoercionMode | None,
+    validate_output: CoercionMode | None,
+) -> Callable[P, R]:
+    signature = inspect.signature(fn)
+
+    @wraps(fn)
+    def wrapper(*args: P.args, **kwds: P.kwargs) -> R:
+        bound_arguments = signature.bind(*args, **kwds)
+        bound_arguments.apply_defaults()
+
+        if validate_input is not None:
+            for parameter_name, expected_schema in spec.input_schemas.items():
+                value = bound_arguments.arguments.get(parameter_name)
+                if not isinstance(value, DataFrame):
+                    raise TypeError(
+                        f"Parameter '{parameter_name}' must be a pyspark.sql.DataFrame"
+                    )
+
+                try:
+                    bound_arguments.arguments[parameter_name] = (
+                        schemas.coerce_dataframe(value, expected_schema, validate_input)
+                    )
+                except ValueError as e:
+                    raise ValueError(
+                        f"Schema mismatch for parameter '{parameter_name}'"
+                    ) from e
+
+        result = fn(*bound_arguments.args, **bound_arguments.kwargs)
+
+        if validate_output is not None and spec.output_schema is not None:
+            if not isinstance(result, DataFrame):
+                raise TypeError(
+                    f"Return value from '{fn.__name__}' must be a pyspark.sql.DataFrame"
+                )
+
+            try:
+                result = cast(
+                    R,
+                    schemas.coerce_dataframe(
+                        result, spec.output_schema, validate_output
+                    ),
+                )
+            except ValueError as e:
+                raise ValueError(f"Return schema mismatch for '{fn.__name__}'") from e
+
+        return result
+
+    return wrapper
+
+
+@overload
+def transform(
+    f: Callable[P, R],
+    *,
+    validate_input: CoercionMode | None = "project_all",
+    validate_output: CoercionMode | None = "project_all",
+) -> Callable[P, R]: ...
+
+
+@overload
+def transform(
+    f: None = None,
+    *,
+    validate_input: CoercionMode | None = "project_all",
+    validate_output: CoercionMode | None = "project_all",
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def transform(
+    f: Callable[P, R] | None = None,
+    *,
+    validate_input: CoercionMode | None = "project_all",
+    validate_output: CoercionMode | None = "project_all",
+):
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        spec = _inspect_transform(fn)
+        return _wrap_transform(
+            fn,
+            spec,
+            validate_input=validate_input,
+            validate_output=validate_output,
+        )
+
+    if f is None:
+        return decorator
+
+    return decorator(f)

spark_joinery-0.0.1/src/spark_joinery/type_inspection.py

@@ -0,0 +1,124 @@
+import datetime
+import decimal
+from types import UnionType
+
+from pyspark.sql import types
+from typing import Any, Annotated, Literal, get_origin, get_args, Union
+
+NoneType = type(None)
+
+
+SPARK_MAX_DECIMAL_PRECISION = 38
+DEFAULT_FRACTIONAL_DIGITS = 18
+
+
+def allows_none(tp: Any) -> bool:
+    origin = get_origin(tp)
+    args = get_args(tp)
+
+    # Optional[T] is Union[T, None]
+    if origin is Union and NoneType in args:
+        return True
+
+    # PEP 604 syntax: T | None
+    if origin is UnionType and NoneType in args:
+        return True
+
+    return False
+
+
+def is_list(tp: Any) -> bool:
+    tp = normalize_python_type(tp)
+    origin = get_origin(tp)
+    return origin is list
+
+
+def get_list_element_type(tp: Any) -> Any:
+    tp = normalize_python_type(tp)
+    if not is_list(tp):
+        raise ValueError(f"{tp} is not a list type")
+    args = get_args(tp)
+    if len(args) != 1:
+        raise ValueError(f"List type {tp} should have exactly one type argument")
+    return args[0]
+
+
+def normalize_python_type(python_type: Any) -> Any:
+    origin = get_origin(python_type)
+    args = get_args(python_type)
+
+    if origin is Annotated and args:
+        return normalize_python_type(args[0])
+
+    if origin in (Union, UnionType) and args:
+        non_none_args = [arg for arg in args if arg is not NoneType]
+        if len(non_none_args) == 1 and len(non_none_args) != len(args):
+            return normalize_python_type(non_none_args[0])
+
+    return python_type
+
+
+def spark_type_from_annotated(python_type: Any) -> types.DataType | None:
+    origin = get_origin(python_type)
+    args = get_args(python_type)
+
+    if origin is not Annotated or len(args) < 2:
+        return None
+
+    for metadata in args[1:]:
+        if isinstance(metadata, types.DataType):
+            return metadata
+
+    return None
+
+
+def spark_type_from_literal(python_type: Any) -> types.DataType | None:
+    if get_origin(python_type) is not Literal:
+        return None
+
+    literal_values = get_args(python_type)
+    if not literal_values:
+        raise ValueError("Literal type must include at least one value")
+
+    spark_types = {
+        get_spark_type_from_python_type(type(literal_value))
+        for literal_value in literal_values
+    }
+
+    if len(spark_types) != 1:
+        raise ValueError(
+            f"Literal values must resolve to a single Spark type: {python_type}"
+        )
+
+    return next(iter(spark_types))
+
+
+def get_spark_type_from_python_type(python_type: type) -> types.DataType:
+    annotated_spark_type = spark_type_from_annotated(python_type)
+    if annotated_spark_type is not None:
+        return annotated_spark_type
+
+    literal_spark_type = spark_type_from_literal(python_type)
+    if literal_spark_type is not None:
+        return literal_spark_type
+
+    python_type = normalize_python_type(python_type)
+
+    if python_type is int:
+        return types.IntegerType()
+    elif python_type is str:
+        return types.StringType()
+    elif python_type is float:
+        return types.FloatType()
+    elif python_type is bool:
+        return types.BooleanType()
+    elif python_type is datetime.datetime:
+        return types.TimestampType()
+    elif python_type is datetime.date:
+        return types.DateType()
+    elif python_type is decimal.Decimal:
+        return types.DecimalType(SPARK_MAX_DECIMAL_PRECISION, DEFAULT_FRACTIONAL_DIGITS)
+    elif python_type is bytes:
+        return types.BinaryType()
+    else:
+        raise ValueError(f"Unsupported Python type: {python_type}")