diff --git a/src/freshdata/learning/diff.py b/src/freshdata/learning/diff.py index a17cf721..643b1636 100644 --- a/src/freshdata/learning/diff.py +++ b/src/freshdata/learning/diff.py @@ -24,7 +24,23 @@ def _diff_kind(raw_missing: bool, clean_missing: bool) -> str: return "value_change" +def _by_value(series: pd.Series) -> pd.Series: + """``series`` with a categorical dtype replaced by its plain object values. + + Comparing two Categorical Series raises ``TypeError`` unless both share + identical categories (and ordering), which a messy/clean pair almost never + does -- repairing the values is exactly what changes the categories. Cells + are compared by value, so decode categoricals to the object values an + object-dtype column would hold. + """ + if isinstance(series.dtype, pd.CategoricalDtype): + return series.astype(object) + return series + + def _column_diffs(messy: pd.Series, clean: pd.Series, column: str) -> list[ValueDiff]: + messy = _by_value(messy) + clean = _by_value(clean) raw_na = messy.isna() clean_na = clean.isna() # A cell differs when values are unequal, except both-missing which is equal. diff --git a/tests/learning/test_categorical.py b/tests/learning/test_categorical.py new file mode 100644 index 00000000..b6f99198 --- /dev/null +++ b/tests/learning/test_categorical.py @@ -0,0 +1,215 @@ +"""``fd.learn`` on categorical columns. + +Comparing two pandas Categoricals raises ``TypeError`` unless their categories +are identical, and a messy/clean pair almost never shares categories -- the +repair is what changes them. Learning must compare, align and learn such +columns by value, exactly as it does for the object-dtype equivalent. The +fixtures avoid version-specific pandas APIs so they run unchanged on pandas +1.5 (Python 3.9 lane) and pandas 2.x. +""" + +from __future__ import annotations + +import json + +import numpy as np +import pandas as pd +import pytest + +import freshdata as fd +from freshdata.learning import LearningProfile, learn, load_profile, save_profile +from freshdata.learning.align import align_pair +from freshdata.learning.diff import compute_diff + +_RAW = [ + "Deliverd", + "SHIPPED", + "shipped ", + "delivered", + "SHIPPED", + "Deliverd", + "shipped ", + "SHIPPED", +] +_CLEAN = [ + "delivered", + "shipped", + "shipped", + "delivered", + "shipped", + "delivered", + "shipped", + "shipped", +] +_IDS = [f"A{i}" for i in range(len(_RAW))] + + +def _frames(raw: list, clean: list) -> tuple[pd.DataFrame, pd.DataFrame]: + ids = [f"A{i}" for i in range(len(raw))] + return ( + pd.DataFrame({"order_id": ids, "status": raw}), + pd.DataFrame({"order_id": ids, "status": clean}), + ) + + +def _as_object(df: pd.DataFrame) -> pd.DataFrame: + """The object-dtype equivalent of ``df`` (categoricals decoded).""" + return df.assign( + **{ + str(c): df[c].astype(object) + for c in df.columns + if isinstance(df[c].dtype, pd.CategoricalDtype) + } + ) + + +def _learned(profile: LearningProfile) -> str: + """Everything a profile learned, as canonical JSON (timestamps excluded).""" + audit = profile.audit().to_dict() + payload = { + "rules": [r.to_dict() for r in profile.rules], + "value_maps": {c: m.to_dict() for c, m in sorted(profile.value_maps.items())}, + "examples": profile.examples.to_dict() if profile.examples is not None else None, + "value_patterns": profile.memory.value_patterns if profile.memory is not None else None, + "notes": audit.get("notes"), + "demotions": audit.get("demotions"), + "holdout_metrics": audit.get("holdout_metrics"), + } + return json.dumps(payload, sort_keys=True, default=str) + + +def _assert_parity(messy: pd.DataFrame, clean: pd.DataFrame, **kwargs) -> LearningProfile: + profile = learn(messy, clean, **kwargs) + reference = learn(_as_object(messy), _as_object(clean), **kwargs) + assert _learned(profile) == _learned(reference) + return profile + + +def test_repro_categoricals_with_different_categories() -> None: + messy = pd.DataFrame({"c": pd.Categorical([" a", "B ", "a"])}) + clean = pd.DataFrame({"c": pd.Categorical(["a", "b", "a"])}) + profile = fd.learn(messy, clean) + assert isinstance(profile, LearningProfile) + assert _learned(profile) == _learned(fd.learn(_as_object(messy), _as_object(clean))) + + +def test_diff_compares_categoricals_by_value() -> None: + messy, clean = _frames(pd.Categorical(_RAW), pd.Categorical(_CLEAN)) + summary = compute_diff(align_pair(messy, clean, key="order_id")) + reference = compute_diff(align_pair(_as_object(messy), _as_object(clean), key="order_id")) + assert {col: [d.to_dict() for d in diffs] for col, diffs in summary.column_diffs.items()} == { + col: [d.to_dict() for d in diffs] for col, diffs in reference.column_diffs.items() + } + pairs = {(d.raw_value, d.clean_value, d.support) for d in summary.column_diffs["status"]} + assert pairs == { + ("SHIPPED", "shipped", 3), + ("Deliverd", "delivered", 2), + ("shipped ", "shipped", 2), + } + # Only the values changed, never the dtype. + assert summary.schema_diffs.dtype_changes == {} + + +@pytest.mark.parametrize( + ("messy_status", "clean_status"), + [ + pytest.param(pd.Categorical(_RAW), pd.Categorical(_CLEAN), id="both-different-categories"), + pytest.param( + pd.Categorical(_RAW, categories=sorted(set(_RAW) | set(_CLEAN))), + pd.Categorical(_CLEAN, categories=sorted(set(_RAW) | set(_CLEAN))), + id="both-same-categories", + ), + pytest.param(pd.Categorical(_RAW), list(_CLEAN), id="categorical-messy-object-clean"), + pytest.param(list(_RAW), pd.Categorical(_CLEAN), id="object-messy-categorical-clean"), + pytest.param( + pd.Categorical(_RAW, ordered=True), + pd.Categorical(_CLEAN, ordered=True), + id="both-ordered", + ), + pytest.param( + pd.Categorical(_RAW, ordered=True), + pd.Categorical(_CLEAN, categories=["shipped", "delivered"], ordered=True), + id="ordered-different-order", + ), + pytest.param( + pd.Categorical(_RAW, ordered=True), pd.Categorical(_CLEAN), id="ordered-vs-unordered" + ), + ], +) +@pytest.mark.parametrize("key", [None, "order_id"]) +def test_learning_matches_object_dtype(messy_status, clean_status, key) -> None: + messy, clean = _frames(messy_status, clean_status) + profile = _assert_parity(messy, clean, key=key, min_support=2) + entries = {(e.raw_value, e.clean_value) for e in profile.value_maps["status"].entries} + assert entries == { + ("SHIPPED", "shipped"), + ("shipped ", "shipped"), + ("Deliverd", "delivered"), + } + + +def test_categorical_key_column() -> None: + messy, clean = _frames(pd.Categorical(_RAW), pd.Categorical(_CLEAN)) + messy["order_id"] = pd.Categorical(messy["order_id"]) + clean = clean.iloc[::-1].reset_index(drop=True) + clean["order_id"] = pd.Categorical(clean["order_id"], categories=list(reversed(_IDS))) + _assert_parity(messy, clean, key="order_id", min_support=2) + + +def test_missing_cells_in_categoricals() -> None: + raw = ["x", "N/A", "y", np.nan, "x", "N/A", "y", np.nan, "x", "N/A"] + clean = ["x", np.nan, "y", np.nan, "x", np.nan, "y", np.nan, "x", np.nan] + messy_df, clean_df = _frames(pd.Categorical(raw), pd.Categorical(clean)) + summary = compute_diff(align_pair(messy_df, clean_df, key="order_id")) + # Both-missing cells are equal; only the sentinel cells differ. + (diff,) = summary.column_diffs["status"] + assert (diff.raw_value, diff.clean_value, diff.support, diff.kind) == ( + "N/A", + None, + 3, + "value_to_missing", + ) + profile = _assert_parity(messy_df, clean_df, key="order_id", min_support=2) + assert any(r.rule == "sentinel" for r in profile.rules) + + +def test_imputed_categorical_cells() -> None: + raw = ["a", np.nan, "a", "b", np.nan, "a", "a", "b", np.nan, "a"] + clean = ["a", "a", "a", "b", "a", "a", "a", "b", "a", "a"] + messy_df, clean_df = _frames(pd.Categorical(raw), pd.Categorical(clean)) + profile = _assert_parity(messy_df, clean_df, key="order_id", min_support=2) + assert profile.value_maps == {} # imputation is never a literal map + + +def test_sensitive_categorical_column_learns_like_object() -> None: + raw = ["A@X.COM ", "b@@y.com", "c@z.com"] * 3 + ["d@@w.com"] + clean = ["a@x.com", "b@y.com", "c@z.com"] * 3 + ["d@w.com"] + messy = pd.DataFrame({"email": pd.Categorical(raw)}) + clean_df = pd.DataFrame({"email": pd.Categorical(clean)}) + _assert_parity(messy, clean_df, privacy="none", include_sensitive=True, min_support=2) + masked = learn(messy, clean_df, min_support=2) + reference = learn(_as_object(messy), _as_object(clean_df), min_support=2) + assert [e.masked for e in masked.value_maps["email"].entries] == [ + e.masked for e in reference.value_maps["email"].entries + ] + + +def test_categorical_profile_round_trips_and_replays_like_object(tmp_path) -> None: + messy, clean = _frames(pd.Categorical(_RAW), pd.Categorical(_CLEAN)) + profile = learn(messy, clean, key="order_id", min_support=2) + reference = learn(_as_object(messy), _as_object(clean), key="order_id", min_support=2) + + path = tmp_path / "categorical.fdprofile" + save_profile(profile, path) + loaded = load_profile(path) + assert loaded.profile_id == profile.profile_id == reference.profile_id + assert _learned(loaded) == _learned(profile) + + batch = pd.DataFrame( + {"order_id": ["B1", "B2", "B3"], "status": ["Deliverd", "SHIPPED", "shipped "]} + ) + categorical_batch = batch.assign(status=pd.Categorical(batch["status"])) + replayed = fd.clean(categorical_batch, profile=loaded, semantic_mode="auto") + expected = fd.clean(batch, profile=reference, semantic_mode="auto") + assert replayed["status"].astype(object).tolist() == ["delivered", "shipped", "shipped"] + assert replayed["status"].astype(object).tolist() == expected["status"].astype(object).tolist()