diff --git a/docs/cleaning-engine.md b/docs/cleaning-engine.md index 1220b81..80a0976 100644 --- a/docs/cleaning-engine.md +++ b/docs/cleaning-engine.md @@ -18,7 +18,7 @@ level, and a confidence score. | order | step | what it does | |---|---|---| | 1 | `column_names` | snake_case names, deduplicate collisions (`"a", "a"` → `"a", "a_2"`) | -| 2 | `strip_whitespace` | trim surrounding whitespace in text cells (internal spacing kept) | +| 2 | `strip_whitespace` | trim surrounding whitespace in text cells (internal spacing kept) — object, `string`, Arrow `string` and categorical columns; a categorical keeps its dtype (its categories are repaired, and ones that become equal merge) | | 3 | `normalize_sentinels` | `"N/A"`, `"null"`, `"-"`, `""`, `"#REF!"`, … → missing | | 4 | `drop_empty_columns` / `drop_empty_rows` | remove all-missing columns and rows | | 5 | `fix_dtypes` | text → numeric (`"$1,234.56"` works) / datetime / boolean, validated | diff --git a/src/freshdata/_util.py b/src/freshdata/_util.py index ba692d3..07c76ef 100644 --- a/src/freshdata/_util.py +++ b/src/freshdata/_util.py @@ -96,7 +96,44 @@ def stringlike_columns(df: pd.DataFrame) -> list: def _is_stringlike_dtype(dtype: object) -> bool: - return pd.api.types.is_object_dtype(dtype) or isinstance(dtype, pd.StringDtype) + return ( + pd.api.types.is_object_dtype(dtype) + or isinstance(dtype, pd.StringDtype) + or is_arrow_string_dtype(dtype) + ) + + +def is_arrow_string_dtype(dtype: object) -> bool: + """True for a ``pd.ArrowDtype`` holding strings (pandas >= 2 only). + + ``pd.ArrowDtype(pa.string())`` carries the same text as ``string[pyarrow]`` + but is a different dtype class, so it needs its own check. pandas 1.5's + experimental ``ArrowDtype`` is left alone. + """ + arrow_dtype_cls = getattr(pd, "ArrowDtype", None) + if PANDAS_MAJOR < 2 or arrow_dtype_cls is None or not isinstance(dtype, arrow_dtype_cls): + return False + import pyarrow as pa # noqa: PLC0415 - an ArrowDtype implies pyarrow is installed + + arrow_type = getattr(dtype, "pyarrow_dtype", None) + is_string_view = getattr(pa.types, "is_string_view", None) + return bool( + pa.types.is_string(arrow_type) + or pa.types.is_large_string(arrow_type) + or (is_string_view is not None and is_string_view(arrow_type)) + ) + + +def as_string_view(s: pd.Series) -> pd.Series: + """``string[pyarrow]`` copy of an Arrow-string column; any other column as-is. + + Type inference parses text into numbers/dates and then does arithmetic on the + result; Arrow-backed results do not implement all of it (e.g. ``%``), while + the ``string[pyarrow]`` path produces regular pandas dtypes. + """ + if is_arrow_string_dtype(s.dtype): + return s.astype(pd.StringDtype("pyarrow")) + return s #: Leading characters Excel/Sheets/LibreOffice interpret as a formula diff --git a/src/freshdata/profile.py b/src/freshdata/profile.py index c8b40a3..cd39ba1 100644 --- a/src/freshdata/profile.py +++ b/src/freshdata/profile.py @@ -19,7 +19,12 @@ from .render.mixins import HtmlReprMixin from .steps.dtypes import suggest_conversion from .steps.outliers import _bounds -from .steps.strings import active_sentinels, normalize_text +from .steps.strings import ( + active_sentinels, + is_text_categorical_dtype, + normalize_categorical, + normalize_text, +) @dataclass(frozen=True) @@ -162,19 +167,22 @@ def _profile_column(name: str, s: pd.Series, config: CleanConfig, issues.append("mixed value types") is_textual = _is_stringlike_dtype(s.dtype) - if is_textual and non_null: - normalized, n_stripped, n_sentinels, n_case = normalize_text(s, config, sentinels) + text_categorical = is_text_categorical_dtype(s.dtype) + if (is_textual or text_categorical) and non_null: + normalize = normalize_categorical if text_categorical else normalize_text + normalized, n_stripped, n_sentinels, n_case = normalize(s, config, sentinels) if n_stripped: issues.append(f"{n_stripped} value(s) with surrounding whitespace") if n_sentinels: issues.append(f"{n_sentinels} sentinel value(s) meaning missing") if n_case: issues.append(f"{n_case} value(s) would be converted to {config.string_case}case") - target, converted, n_coerced = suggest_conversion(normalized, config) - if converted is not None: - suggested = str(converted.dtype) - note = f", {n_coerced} unparseable" if n_coerced else "" - issues.append(f"would convert to {suggested}{note}") + if is_textual: # cleaning keeps categoricals categorical; no dtype suggestion + target, converted, n_coerced = suggest_conversion(normalized, config) + if converted is not None: + suggested = str(converted.dtype) + note = f", {n_coerced} unparseable" if n_coerced else "" + issues.append(f"would convert to {suggested}{note}") if is_numeric_dtype(s) and not is_bool_dtype(s) and non_null >= 20: bounds = _bounds(s, config) diff --git a/src/freshdata/steps/dtypes.py b/src/freshdata/steps/dtypes.py index 795771b..a300575 100644 --- a/src/freshdata/steps/dtypes.py +++ b/src/freshdata/steps/dtypes.py @@ -23,7 +23,13 @@ import pandas as pd from pandas.api.types import infer_dtype, is_datetime64_any_dtype -from .._util import PANDAS_MAJOR, mask_sensitive_value, sample_series, stringlike_columns +from .._util import ( + PANDAS_MAJOR, + as_string_view, + mask_sensitive_value, + sample_series, + stringlike_columns, +) from ..config import CleanConfig from ..report import CleanReport @@ -458,6 +464,7 @@ def suggest_conversion( the cleaning pipeline and :func:`freshdata.profile` so the preview always matches what cleaning would actually do. """ + s = as_string_view(s) # Arrow strings parse through the string[pyarrow] path nonnull = s.dropna() if nonnull.empty: return "none", None, 0 @@ -665,14 +672,15 @@ def fix_dtypes(df: pd.DataFrame, config: CleanConfig, report: CleanReport) -> pd for col in stringlike_columns(df): if str(col) in protected: continue # context-protected columns must stay byte-identical - target, converted, n_coerced = suggest_conversion(df[col], config) + s = as_string_view(df[col]) + target, converted, n_coerced = suggest_conversion(s, config) if converted is None: - _warn_type_contamination(str(col), df[col], config, report) + _warn_type_contamination(str(col), s, config, report) continue description = f"converted to {converted.dtype}" if n_coerced: description += f" ({n_coerced} unparseable value(s) set to missing)" - _record_coerced(str(col), df[col], converted, report, config) + _record_coerced(str(col), s, converted, report, config) report.add("fix_dtypes", description, column=str(col), count=int(converted.notna().sum()) + n_coerced) df[col] = converted diff --git a/src/freshdata/steps/strings.py b/src/freshdata/steps/strings.py index d6dcf4c..0744417 100644 --- a/src/freshdata/steps/strings.py +++ b/src/freshdata/steps/strings.py @@ -15,7 +15,7 @@ from pandas.api.types import infer_dtype from .._sentinels import DEFAULT_SENTINELS -from .._util import stringlike_columns +from .._util import _is_stringlike_dtype from ..config import CleanConfig from ..report import CleanReport @@ -91,8 +91,51 @@ def normalize_text( return s, n_stripped, n_sentinels, n_case +def is_text_categorical_dtype(dtype: object) -> bool: + """True for a categorical dtype whose categories hold text worth repairing.""" + return ( + isinstance(dtype, pd.CategoricalDtype) + and infer_dtype(dtype.categories, skipna=True) in _TEXTUAL_KINDS + ) + + +def normalize_categorical( + s: pd.Series, config: CleanConfig, sentinels: frozenset[str] +) -> tuple[pd.Series, int, int, int]: + """:func:`normalize_text` for a text categorical, keeping the categorical dtype. + + Values are repaired exactly as the equivalent object column would be (so + counts match), then rebuilt as a categorical with the same ``ordered`` flag + whose categories are the repaired originals: ``" a "`` and ``"a"`` merge, + and sentinel categories disappear. + """ + normalized, n_stripped, n_sentinels, n_case = normalize_text( + s.astype(object), config, sentinels + ) + if not (n_stripped or n_sentinels or n_case): + return s, 0, 0, 0 + categories, *_ = normalize_text(pd.Series(s.cat.categories, dtype=object), config, sentinels) + rebuilt = pd.Categorical( + normalized, categories=pd.unique(categories.dropna()), ordered=s.cat.ordered + ) + return pd.Series(rebuilt, index=s.index, name=s.name), n_stripped, n_sentinels, n_case + + +def _text_columns(df: pd.DataFrame) -> list: + """Object/string columns plus text categoricals, in frame order.""" + return [ + col + for col, dtype in zip(df.columns, df.dtypes) + if _is_stringlike_dtype(dtype) or is_text_categorical_dtype(dtype) + ] + + def clean_strings(df: pd.DataFrame, config: CleanConfig, report: CleanReport) -> pd.DataFrame: - """Apply whitespace stripping and sentinel→missing to text-capable columns.""" + """Apply whitespace stripping and sentinel→missing to text-capable columns. + + Categorical columns with text categories are repaired too; they keep their + categorical dtype (see :func:`normalize_categorical`). + """ if not ( config.strip_whitespace or config.normalize_sentinels @@ -103,10 +146,14 @@ def clean_strings(df: pd.DataFrame, config: CleanConfig, report: CleanReport) -> from ..guard import hard_protected_columns # noqa: PLC0415 — cycle-safe lazy import protected = hard_protected_columns(config, df.columns) - for col in stringlike_columns(df): + for col in _text_columns(df): if str(col) in protected: continue # context-protected columns must stay byte-identical - normalized, n_stripped, n_sentinels, n_case = normalize_text(df[col], config, sentinels) + s = df[col] + normalize = ( + normalize_categorical if isinstance(s.dtype, pd.CategoricalDtype) else normalize_text + ) + normalized, n_stripped, n_sentinels, n_case = normalize(s, config, sentinels) if n_stripped: report.add("strip_whitespace", "trimmed surrounding whitespace", column=str(col), count=n_stripped) diff --git a/src/freshdata/textclean.py b/src/freshdata/textclean.py index 9ff1726..d0379a9 100644 --- a/src/freshdata/textclean.py +++ b/src/freshdata/textclean.py @@ -28,6 +28,8 @@ import pandas as pd +from ._util import _is_stringlike_dtype + __all__ = [ "TextCleanConfig", "CleanedText", @@ -301,7 +303,7 @@ def clean_text( never runs on an amount or identifier column. """ if columns is None: - cols = [c for c in df.columns if df[c].dtype == object or str(df[c].dtype) == "string"] + cols = [c for c in df.columns if _is_stringlike_dtype(df[c].dtype)] else: missing = [c for c in columns if c not in df.columns] if missing: diff --git a/tests/test_profile.py b/tests/test_profile.py index 6b34ef2..2be5ede 100644 --- a/tests/test_profile.py +++ b/tests/test_profile.py @@ -96,3 +96,21 @@ def test_empty_frame_profile(): p = fd.profile(pd.DataFrame()) assert p.n_rows == 0 and p.n_cols == 0 assert str(p) # renders without crashing + + +def test_profile_flags_text_issues_in_categorical_columns(): + df = pd.DataFrame({"c": pd.Categorical([" a ", "N/A", "b", "b"])}) + issues = fd.profile(df).columns[0].issues + assert "1 value(s) with surrounding whitespace" in issues + assert "1 sentinel value(s) meaning missing" in issues + assert not any("would convert" in issue for issue in issues) # stays categorical + + +def test_profile_flags_text_issues_in_arrow_string_columns(): + if int(pd.__version__.split(".")[0]) < 2: + pytest.skip("pd.ArrowDtype strings need pandas >= 2") + pa = pytest.importorskip("pyarrow") + df = pd.DataFrame({"s": pd.Series([" a ", "N/A", "b", "b"], dtype=pd.ArrowDtype(pa.string()))}) + issues = fd.profile(df).columns[0].issues + assert "1 value(s) with surrounding whitespace" in issues + assert "1 sentinel value(s) meaning missing" in issues diff --git a/tests/test_strings.py b/tests/test_strings.py index 33fc84b..f6d9c7f 100644 --- a/tests/test_strings.py +++ b/tests/test_strings.py @@ -1,7 +1,9 @@ import numpy as np import pandas as pd +import pytest import freshdata as fd +from freshdata._util import PANDAS_MAJOR def test_whitespace_stripped_object_and_string_dtype(): @@ -88,3 +90,62 @@ def test_unhashable_values_pass_through(): out = fd.clean(df) assert out["v"].iloc[0] == [1, 2] assert not np.any(out["w"].isna()) + + +def _plain(values): + return [None if pd.isna(v) else v for v in values] + + +def _column_steps(report, column): + return [(a.step, a.count) for a in report if a.column == column] + + +@pytest.mark.skipif(PANDAS_MAJOR < 2, reason="pd.ArrowDtype strings need pandas >= 2") +@pytest.mark.parametrize( + "values", + [ + [" a ", "N/A", "3", "4"], # text: strip + sentinel + ["1", " 2 ", "N/A", "4"], # numeric-looking: fix_dtypes converts it + ["2024-01-01", " 2024-02-01", None, "2024-03-01"], # dates + ], +) +def test_arrow_string_column_cleans_like_string_pyarrow(values): + pa = pytest.importorskip("pyarrow") + arrow = pd.DataFrame( + {"s": pd.Series(values, dtype=pd.ArrowDtype(pa.string())), "k": [1.0, 2.0, 3.0, 4.0]} + ) + string = arrow.astype({"s": "string[pyarrow]"}) + out_arrow, report_arrow = fd.clean(arrow, return_report=True, verbose=False) + out_string, report_string = fd.clean(string, return_report=True, verbose=False) + assert _plain(out_arrow["s"].astype(object)) == _plain(out_string["s"].astype(object)) + assert _column_steps(report_arrow, "s") == _column_steps(report_string, "s") + assert ("strip_whitespace", 1) in _column_steps(report_arrow, "s") + + +def test_categorical_text_is_normalized_and_stays_categorical(): + cat = pd.Categorical( + [" a ", "N/A", "b", "null", "a"], + categories=["a", " a ", "N/A", "b", "null"], + ordered=True, + ) + df = pd.DataFrame({"c": cat, "k": [1.0, 2.0, 3.0, 4.0, 5.0]}) + out, report = fd.clean( + df, strategy="conservative", return_report=True, verbose=False, + drop_empty_rows=False, drop_duplicates=False, + ) + assert isinstance(out["c"].dtype, pd.CategoricalDtype) + assert out["c"].cat.ordered + assert list(out["c"].cat.categories) == ["a", "b"] # " a " merged, sentinels gone + assert _plain(out["c"]) == ["a", None, "b", None, "a"] + counts = dict(_column_steps(report, "c")) + assert counts["strip_whitespace"] == 1 + assert counts["normalize_sentinels"] == 2 + + +def test_categorical_values_match_object_column(): + values = [" a ", "N/A", "b", "null"] + cat = pd.DataFrame({"c": pd.Categorical(values), "k": [1.0, 2.0, 3.0, 4.0]}) + out_cat = fd.clean(cat, verbose=False) + out_obj = fd.clean(cat.astype({"c": object}), verbose=False) + assert isinstance(out_cat["c"].dtype, pd.CategoricalDtype) + assert _plain(out_cat["c"].astype(object)) == _plain(out_obj["c"])