Data Contracts, Quality, and Features

The most expensive ML bugs are not model bugs. They are data bugs that produce models which look correct offline and fail silently in production. A model trained on corrupted or drifted features does not throw an exception. It makes confident wrong predictions until someone notices the business metric moving in the wrong direction three weeks later.

Data contracts are the mechanism that prevents this. They encode your expectations about incoming data and fail loudly when those expectations are violated.

What Is a Data Contract?

A data contract is a formal specification of what a dataset or feature is expected to look like: its schema (column names and types), its statistical properties (value ranges, null rates, distributions), and its business rules (no negative prices, user IDs must be in the user table).

The key insight is that ML pipelines have two failure modes: model failure (the model does not generalize) and data failure (the model never received correct inputs). Data contracts catch the second type.

Defining Expectations with Pandera

Pandera lets you define DataFrame schemas with statistical validation in Python code that lives in your repository.

pythonCopy
import pandera as pa
from pandera import Column, DataFrameSchema, Check
import pandas as pd
import numpy as np

# Define the contract for your feature table
feature_schema = DataFrameSchema(
    columns={
        'user_id': Column(str, nullable=False),
        'age': Column(float, checks=[
            Check.ge(0),     # >= 0
            Check.le(120),   # <= 120
        ], nullable=True),
        'days_since_last_purchase': Column(float, checks=[
            Check.ge(0),
        ], nullable=False),
        'total_spend_30d': Column(float, checks=[
            Check.ge(0),
            Check.le(1e7),  # sanity cap
        ], nullable=False),
        'label': Column(int, checks=[Check.isin([0, 1])], nullable=False),
    },
    checks=[
        # At most 5% null rate on age
        Check(lambda df: df['age'].isnull().mean() < 0.05,
              error="age null rate exceeds 5%"),
        # At least 1000 rows
        Check(lambda df: len(df) >= 1000,
              error="dataset too small"),
    ]
)

# Validate at the pipeline entry point
def load_and_validate_features(path: str) -> pd.DataFrame:
    df = pd.read_parquet(path)
    return feature_schema.validate(df)  # raises SchemaError if violated

If the contract is violated, SchemaError is raised immediately with a clear message. Your pipeline fails fast instead of silently propagating bad data.

Detecting Distribution Drift with Great Expectations

Great Expectations goes further than schema validation - it can detect when the statistical distribution of a feature shifts over time.

pythonCopy
import great_expectations as gx

context = gx.get_context()

# Define expectation suite
suite = context.add_expectation_suite("feature_table")

validator = context.get_validator(
    datasource_name="my_datasource",
    data_connector_name="default_inferred_data_connector_name",
    data_asset_name="features.parquet",
    expectation_suite=suite
)

# Column existence and types
validator.expect_column_to_exist("user_id")
validator.expect_column_values_to_not_be_null("user_id")

# Statistical range expectations
validator.expect_column_values_to_be_between(
    "age", min_value=0, max_value=120, mostly=0.99  # 99% must pass
)

# Null rate threshold
validator.expect_column_values_to_not_be_null("total_spend_30d")

# Distribution expectation: mean within range
validator.expect_column_mean_to_be_between(
    "days_since_last_purchase", min_value=5, max_value=60
)

validator.save_expectation_suite()

Run this expectation suite on every new batch of data before training or serving. The HTML report tells you exactly which expectations failed and by how much.

Fail Fast vs. Alert: Choosing the Right Response

Not every contract violation should stop your pipeline. You need a severity model:

pythonCopy
from enum import Enum
from typing import Optional

class ViolationSeverity(Enum):
    CRITICAL = "critical"   # halt pipeline
    WARNING = "warning"     # alert and continue
    INFO = "info"           # log only

def check_feature_quality(df: pd.DataFrame) -> list[dict]:
    violations = []

    # Critical: missing user_id means broken join
    null_user_ids = df['user_id'].isnull().sum()
    if null_user_ids > 0:
        violations.append({
            'column': 'user_id',
            'issue': f'{null_user_ids} null user IDs',
            'severity': ViolationSeverity.CRITICAL
        })

    # Warning: age null rate slightly elevated
    age_null_rate = df['age'].isnull().mean()
    if 0.05 <= age_null_rate < 0.20:
        violations.append({
            'column': 'age',
            'issue': f'age null rate {age_null_rate:.1%}',
            'severity': ViolationSeverity.WARNING
        })

    # Critical: age null rate too high (likely broken join)
    if age_null_rate >= 0.20:
        violations.append({
            'column': 'age',
            'issue': f'age null rate {age_null_rate:.1%}  -  possible broken join',
            'severity': ViolationSeverity.CRITICAL
        })

    return violations

def run_pipeline(df: pd.DataFrame):
    violations = check_feature_quality(df)
    critical = [v for v in violations if v['severity'] == ViolationSeverity.CRITICAL]

    if critical:
        raise ValueError(f"Critical data quality violations: {critical}")

    warnings = [v for v in violations if v['severity'] == ViolationSeverity.WARNING]
    for w in warnings:
        print(f"WARNING: {w['column']}  -  {w['issue']}")

The Feature Store Connection

Data contracts become even more valuable when you have a feature store. The contract is the interface between the team that produces features and the team that trains models. It decouples producers from consumers: as long as the contract is met, the upstream can change how features are computed.

Common Mistakes

Validating after transformation: Contracts should run on raw inputs, not on the transformed output. If your transformation is the thing that broke the data, post-transformation validation will miss it.

Expectations that are too loose: A null rate expectation of "less than 100%" is useless. Set thresholds based on historical data from healthy production batches, not arbitrary round numbers.

No monitoring in serving: Training-time contracts are not enough. Run the same expectations on the features your serving system receives. Training/serving skew is a common production failure.

Where to Go Next

• leakage-proof-feature-pipelines - quality data that still leaks the target will fool you
• python-ml-bridge-for-swe - the pandas and numpy patterns for efficient feature validation
• training-pipelines-experiment-strategy - integrate data contracts into your training pipeline