fix(text): normalize ArrowDtype string and categorical columns

Whitespace stripping, sentinel normalization, fix_dtypes and the profile's
text issues skipped two kinds of text column:

- pd.ArrowDtype(pa.string()) (pandas >= 2): _is_stringlike_dtype only knew
  object and StringDtype. It now accepts Arrow string/large_string/string_view.
  Type inference parses through a string[pyarrow] view, because Arrow-backed
  parse results lack arithmetic the numeric check relies on (e.g. `%`), so the
  column cleans exactly like the same data as string[pyarrow].
- Categorical columns: categories with text are repaired too, and the column
  keeps its categorical dtype and `ordered` flag. Values are normalized like
  the equivalent object column (same counts); categories that become equal
  merge and sentinel categories disappear. fix_dtypes still leaves categoricals
  alone, and the profile reports their whitespace/sentinel issues without a
  dtype suggestion.

fd.clean_text's default column selection uses the same text-dtype check.

Closes #212

Author: kevincostner17

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 |

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

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)

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

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)

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:

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

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"])