pyspark-contracts

Lightweight schema and data quality validation for PySpark DataFrames.

Validates at runtime — before your pipeline runs — with a structured JSON report queryable in CloudWatch Logs Insights, Datadog, or any log aggregator.

Install

pip install pyspark-contracts

PySpark is assumed to be on your classpath (Glue, Databricks, EMR). For local development: pip install pyspark-contracts[pyspark]

Quick start

from pyspark.sql.types import DateType, FloatType, StringType
from pyspark_contracts import Contract, Field

class OdometerContract(Contract):
    vin           = Field(StringType(), nullable=False, min_length=17, max_length=17)
    odometerStart = Field(FloatType(),  nullable=False, min_value=0)
    odometerEnd   = Field(FloatType(),  nullable=False, min_value=0)
    readingDate   = Field(DateType())

# Hard fail (default) — logs ERROR + raises ContractViolationError
OdometerContract().validate(df)

# Soft fail — logs WARNING + returns ViolationReport
report = OdometerContract().validate(df, mode="soft")
if report:
    print(f"{len(report.violations)} violation(s) found")

Validation modes

Lazy vs fail-fast

By default, "hard" mode stops at the first violation found (fewer Spark actions, faster feedback), while "soft" mode collects every violation so you get the full picture in one run. Override either with lazy:

# Hard mode, but collect every violation before raising
OdometerContract().validate(df, lazy=True)

# Soft mode, but stop at the first violation
report = OdometerContract().validate(df, mode="soft", lazy=False)

Cross-column checks

Field constraints are column-local. For rules spanning more than one column, use condition (simple two-column comparisons) or @check (arbitrary logic):

from pyspark_contracts import Contract, Field, check
from pyspark.sql import DataFrame
from pyspark.sql import functions as F
from pyspark.sql.types import DoubleType, TimestampType

class FinalDfContract(Contract):
    end_dt   = Field(TimestampType(), nullable=False)
    start_dt = Field(TimestampType(), nullable=False,
                     condition=lambda f: F.col(f) < F.col("end_dt"),
                     condition_description="start_dt must precede end_dt")

    delta_km = Field(DoubleType(), nullable=False)

    @check("delta_km must be non-negative (no regression)")
    def no_regression(self, df: DataFrame) -> DataFrame:
        return df.filter(F.col("delta_km") < 0)

    @check("delta_km must not exceed daily threshold")
    def no_abnormal_jump(self, df: DataFrame, threshold: float = 1000.0) -> DataFrame:
        return df.filter(F.col("delta_km") > threshold)

A @check method receives the full DataFrame and returns the failing rows — an empty result means it passed. Extra parameters (threshold) are supplied via validate() and routed to the checks that declare them:

report = FinalDfContract().validate(df, mode="soft", threshold=2000.0)

Both condition and @check are skipped entirely when the DataFrame has a missing column or type mismatch, since either can reference arbitrary columns.

Validating pipeline functions

@check_input and @check_output move contract validation from the function body onto its signature, so a refactor can't accidentally drop it:

from pyspark_contracts import check_input, check_output

@check_input(AccumulatedDataContract, param="accumulated_df")
@check_output(FinalDfContract)
def build_final_df(accumulated_df: DataFrame, contracts_df: DataFrame) -> DataFrame:
    return accumulated_df.join(contracts_df, ...)

check_input validates the named parameter before the function runs — a bad step never executes. check_output validates the return value after the function runs, before it's returned to the caller. Both accept mode and any extra kwargs, forwarded to validate() exactly as if you'd called it directly:

@check_output(FinalDfContract, mode="soft", threshold=2000.0)
def build_final_df(df: DataFrame) -> DataFrame:
    ...

A function needing more than one input validated stacks @check_input:

@check_input(AccumulatedDataContract, param="accumulated_df")
@check_input(ContractsContract, param="contracts_df")
def build_final_df(accumulated_df, contracts_df): ...

Validation depth

from pyspark_contracts import ValidationDepth

# Column names and types only — reads df.schema, triggers zero Spark actions
OdometerContract().validate(df, depth=ValidationDepth.SCHEMA_ONLY)

# Skip structural checks, run only nullable/range/length/regex/condition/@check
OdometerContract().validate(df, depth=ValidationDepth.DATA_ONLY)

# Default: both
OdometerContract().validate(df, depth=ValidationDepth.SCHEMA_AND_DATA)

SCHEMA_AND_DATA is the default and runs the full pipeline described everywhere else in this README — every schema and quality check, exactly as if depth were never passed.

SCHEMA_ONLY is particularly useful in unit tests where the DataFrame is built from a known factory: the schema is guaranteed by construction, so only the data constraints are worth checking — the reverse case, DATA_ONLY, skips column/type checks and assumes the schema is already right, without adding a safety net if it isn't.

Disabling validation

PYSPARK_CONTRACTS_ENABLED=false  # skips all validation, no code changes needed

Useful when performance is critical and the schema is already trusted — e.g. the read right after a validated write.

Bringing your own logger

import logging
logger = logging.getLogger("my-etl-job")
OdometerContract().validate(df, logger=logger)

If no logger is passed, the library emits JSON to stdout.

Field options

Documenting a contract

description and metadata on Field are pure documentation — they never affect validation:

class AssetHdrTelemetryReadingContract(Contract):
    asset_hdr_id = Field(LongType(), nullable=False,
                         description="AF internal vehicle identifier")
    meter_reading_value = Field(DecimalType(26, 12), nullable=False, min_value=0,
                                description="Odometer reading in km",
                                metadata={"unit": "km", "source": "asset_hdr_telemetry_reading"})

Export the whole contract as data — types, constraints, descriptions, and the name/description of every @check (never the executable code, which can't be serialized):

schema_dict = AssetHdrTelemetryReadingContract.to_dict()
json_str = AssetHdrTelemetryReadingContract.to_json()

Or print a quick human-readable summary:

print(AssetHdrTelemetryReadingContract.describe())

Log output

Both modes emit the same structured JSON before acting:

{
  "level": "ERROR",
  "message": "contract violation — job aborted",
  "contract": "OdometerContract",
  "mode": "hard",
  "depth": "schema_and_data",
  "violations": [
    {"kind": "missing_column", "column": "odometerStart", "expected_type": "FloatType"},
    {
      "kind": "value_out_of_range",
      "column": "odometerEnd",
      "constraint": "min_value=0",
      "row_pct": 2.3,
      "failure_count": 42,
      "sample_values": [-1.0, -5.5, -12.0]
    }
  ],
  "row_count": 1842
}

CloudWatch Logs Insights

-- All contract violations in the last 7 days
filter message = "contract violation — job aborted"
| stats count(*) as failures by contract

-- Violations by column across all runs
filter ispresent(violations)
| flatten violations
| stats count(*) by violations.kind, violations.column

Violation kinds

Performance note

Each quality constraint triggers a Spark action (filter + count). For DataFrames with many constraints, consider caching before validation:

df = df.cache()
OdometerContract().validate(df)