S053 · Data Quality — Freshness, Volume, Schema, Distribution

"""
dq_demo.py — Four canonical checks on a local Parquet ‘orders’ table.
 
Run:
    pip install pandas pyarrow
    python dq_demo.py
"""
from __future__ import annotations
 
import json
import sys
from dataclasses import asdict, dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Callable
 
import pandas as pd
 
# ---------- fixture: a fake day-of-orders table ----------
TABLE = Path("/tmp/dq_demo_orders.parquet")
 
 
def build_fixture(bad: bool = False) -> pd.DataFrame:
    now = datetime.now(timezone.utc)
    rows = [
        {"order_id": i, "amount": (i % 50) + 1.0, "customer": f"c{i % 200}",
         "ingested_at": now - timedelta(minutes=i % 60)}
        for i in range(5000)
    ]
    df = pd.DataFrame(rows)
    if bad:
        # Simulate three regressions: stale, thin, and skewed.
        df = df.head(500)                                      # volume drop
        df["ingested_at"] = now - timedelta(hours=6)           # stale
        df.loc[df.index[:100], "amount"] = -1                  # negative $$$
    df.to_parquet(TABLE)
    return df
 
 
# ---------- check framework (30 lines is all you need) ----------
@dataclass
class CheckResult:
    check_id: str
    passed: bool
    value: float | str | None
    expected: str
    severity: str  # "critical" | "warn" | "info"
    run_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
 
 
def freshness(df: pd.DataFrame, max_lag_min: int = 120) -> CheckResult:
    lag_min = (datetime.now(timezone.utc) - df["ingested_at"].max()).total_seconds() / 60
    return CheckResult(
        check_id="freshness.orders.ingested_at",
        passed=lag_min <= max_lag_min,
        value=round(lag_min, 1),
        expected=f"<= {max_lag_min} min",
        severity="critical",
    )
 
 
def volume(df: pd.DataFrame, expected: int = 5000, tolerance: float = 0.3) -> CheckResult:
    n = len(df)
    lo, hi = expected * (1 - tolerance), expected * (1 + tolerance)
    return CheckResult(
        check_id="volume.orders.rowcount",
        passed=lo <= n <= hi,
        value=n,
        expected=f"between {int(lo)} and {int(hi)}",
        severity="critical",
    )
 
 
def schema(df: pd.DataFrame, contract: dict[str, str]) -> CheckResult:
    actual = {c: str(t) for c, t in df.dtypes.items()}
    missing = [c for c in contract if c not in actual]
    wrong_type = [c for c in contract if c in actual and not actual[c].startswith(contract[c])]
    ok = not missing and not wrong_type
    return CheckResult(
        check_id="schema.orders.contract_v3",
        passed=ok,
        value=json.dumps({"missing": missing, "wrong_type": wrong_type}),
        expected=json.dumps(contract),
        severity="critical",
    )
 
 
def distribution_amount(df: pd.DataFrame) -> CheckResult:
    neg_rate = (df["amount"] < 0).mean()
    return CheckResult(
        check_id="distribution.orders.amount.non_negative",
        passed=neg_rate == 0,
        value=round(neg_rate, 4),
        expected="0.0",
        severity="warn",
    )
 
 
CHECKS: list[Callable[[pd.DataFrame], CheckResult]] = [
    freshness,
    volume,
    lambda df: schema(df, {"order_id": "int", "amount": "float", "customer": "object", "ingested_at": "datetime"}),
    distribution_amount,
]
 
 
# ---------- runner ----------
def run(df: pd.DataFrame) -> list[CheckResult]:
    results = [c(df) for c in CHECKS]
    for r in results:
        icon = "✅" if r.passed else "❌"
        print(f"{icon} [{r.severity:>8}] {r.check_id}: value={r.value} expected={r.expected}")
    return results
 
 
def route(results: list[CheckResult]) -> int:
    """Return exit code: 0 clean, 1 warn, 2 critical."""
    failed = [r for r in results if not r.passed]
    if any(r.severity == "critical" for r in failed):
        print("\n🚨 PAGE on-call — critical DQ failure. Halting downstream.")
        return 2
    if failed:
        print("\n⚠️  Ticket filed — non-blocking warn.")
        return 1
    print("\n🟢 All checks green.")
    return 0
 
 
def main() -> None:
    mode = sys.argv[1] if len(sys.argv) > 1 else "good"
    df = build_fixture(bad=(mode == "bad"))
    exit_code = route(run(df))
    sys.exit(exit_code)
 
 
if __name__ == "__main__":
    main()

python dq_demo.py good   # exits 0, all green
python dq_demo.py bad    # exits 2, three failures, on-call paged