Session 2.2: Validating Data

# Import os to read env vars
import os
# Import bob and helper functions
from bia_bob import bob, available_models, fix, doc

# Initialize assistant, reads API key from env var automatically
bob.initialize(
    # read endpoint URL from env var
    endpoint=os.getenv('ENDPOINT_URL'),
    # select an available model for your purpose
    model="coder",
    # read the Data Steward system prompt from env var
    system_prompt=os.getenv('SYSTEM_PROMPT_DATA_STEWARD'),
)

import pandas as pd

df = pd.read_csv("../data/field_samples_raw.csv")
df.dtypes

#!uv add pandera

import pandera.pandas as pa
from pandera.pandas import Column, DataFrameSchema, Check

%%bob
Write a pandera DataFrameSchema (using `import pandera.pandas as pa`) for a pandas DataFrame
with these columns and rules:
- sample_id: string, unique, not null, must match the pattern S followed by 4 digits
- species: string, must be one of: Adelie, Chinstrap, Gentoo
- island: string, must be one of: Biscoe, Dream, Torgersen
- date_collected: string, not null, must match an ISO date pattern YYYY-MM-DD
- body_mass_g: integer (coerce), not null, between 2500 and 6500
- flipper_length_mm: float (coerce), not null, between 150 and 250
- sex: string, nullable, must be one of: male, female
- region: string, must equal "Antarctica"
- notes: string, nullable, not required
Store it in a variable called `my_pandera_schema`. Just the code.

schema_reference = DataFrameSchema(
    columns={
        # unique ID, no whitespace, exact pattern "S" + 4 digits
        "sample_id": Column(
            str, unique=True, nullable=False,
            checks=Check.str_matches(r"^S\d{4}$"),
            description="Unique sample identifier, pattern S + 4 digits (e.g. S0001).",
        ),
        # controlled vocabulary -> only these three values are valid
        "species": Column(
            str, checks=Check.isin(["Adelie", "Chinstrap", "Gentoo"]),
            description="Penguin species (controlled vocabulary).",
        ),
        "island": Column(
            str, checks=Check.isin(["Biscoe", "Dream", "Torgersen"]),
            description="Island in the Palmer Archipelago (controlled vocabulary).",
        ),
        # date must be written as ISO 8601 (YYYY-MM-DD)
        "date_collected": Column(
            str, nullable=False,
            checks=Check.str_matches(r"^\d{4}-\d{2}-\d{2}$"),
            description="Date the sample was collected, ISO 8601 (YYYY-MM-DD).",
        ),
        # numbers: coerce text -> number, then check a plausible range
        "body_mass_g": Column(
            int, coerce=True, nullable=False,
            checks=Check.in_range(2500, 6500),
            description="Body mass in grams.",
        ),
        "flipper_length_mm": Column(
            float, coerce=True, nullable=False,
            checks=Check.in_range(150, 250),
            description="Flipper length in millimetres.",
        ),
        # sex may be empty (unknown), but if present must be male/female
        "sex": Column(
            str, nullable=True, checks=Check.isin(["male", "female"]),
            description="Recorded sex; empty means unknown/not recorded.",
        ),
        # constant column
        "region": Column(
            str, checks=Check.eq("Antarctica"),
            description="Broad geographic region (constant across the dataset).",
        ),
        # free-text, optional column
        "notes": Column(
            str, nullable=True, required=False,
            description="Optional free-text field-collection remarks.",
        ),
    },
    strict=False,  # allow extra columns we did not list
)

schema_reference

%%bob
Tell me wether the content of the two variables {my_pandera_schema} and {schema_reference} 
is equal. No code, just a short answer. Point out any differences

try:
    schema_reference.validate(df, lazy=True)
    print("✅ All records are valid.")
except pa.errors.SchemaErrors as exc:
    failures = exc.failure_cases
    print(f"❌ Found {len(failures)} failing cases.\n")
    # The most useful steward's view: how many failures per column and per rule
    summary = failures.groupby(["column", "check"]).size()
    print(summary.to_string())

try:
    schema_reference.validate(df, lazy=True)
except pa.errors.SchemaErrors as exc:
    failures = exc.failure_cases
    for col in ["species", "body_mass_g"]:
        bad_values = failures.loc[failures["column"] == col, "failure_case"].dropna().unique()
        print(f"{col} - distinct invalid values:")
        print(" ", list(bad_values)[:12])
        print()

from pathlib import Path

# Make sure an outputs folder exists, then save the schema as a shareable YAML file
out_dir = Path("../outputs")
out_dir.mkdir(exist_ok=True)

schema_reference.to_yaml(out_dir / "data_schema.yml")
print(f"Saved schema to {out_dir / 'data_schema.yml'}")

%%bob
Show me one extra pandera Check I can add to the body_mass_g column 
in {schema_reference} to reject values that are zero or negative. 
Explain in one sentence what it does.