Repository: capitalone/datacompy
Pull request #415 (2025-05-27 00:23:08)

chore: normalize_string_column for polars and pandas
- adding in a util function `normalize_string_column` for both pandas and polars to reduce redundant code.
- some general cleanup of docstrings
- additional test cases for categoricals.



Patch:
diff --git a/datacompy/core.py b/datacompy/core.py
index 09ea181..64c80e2 100644
--- a/datacompy/core.py
+++ b/datacompy/core.py
@@ -66,9 +66,9 @@ class Compare(BaseCompare):
         A string name for the second dataframe
     ignore_spaces : bool, optional
         Flag to strip whitespace (including newlines) from string columns (including any join
-        columns)
+        columns). Excludes categoricals.
     ignore_case : bool, optional
-        Flag to ignore the case of string columns
+        Flag to ignore the case of string columns. Excludes categoricals.
     cast_column_names_lower: bool, optional
         Boolean indicator that controls of column names will be cast into lower case
 
@@ -285,16 +285,13 @@ class Compare(BaseCompare):
         else:
             params = {"on": self.join_columns}
 
-        if ignore_spaces:
-            for column in self.join_columns:
-                if self.df1[column].dtype.kind == "O" and pd.api.types.is_string_dtype(
-                    self.df1[column]
-                ):
-                    self.df1[column] = self.df1[column].str.strip()
-                if self.df2[column].dtype.kind == "O" and pd.api.types.is_string_dtype(
-                    self.df2[column]
-                ):
-                    self.df2[column] = self.df2[column].str.strip()
+        for column in self.join_columns:
+            self.df1[column] = normalize_string_column(
+                self.df1[column], ignore_spaces=ignore_spaces, ignore_case=False
+            )
+            self.df2[column] = normalize_string_column(
+                self.df2[column], ignore_spaces=ignore_spaces, ignore_case=False
+            )
 
         outer_join = self.df1.merge(
             self.df2,
@@ -870,16 +867,13 @@ def columns_equal(
     """
     default_value = "DATACOMPY_NULL"
     compare: pd.Series[bool]
-    if ignore_spaces:
-        if col_1.dtype.kind == "O" and pd.api.types.is_string_dtype(col_1):
-            col_1 = col_1.str.strip()
-        if col_2.dtype.kind == "O" and pd.api.types.is_string_dtype(col_2):
-            col_2 = col_2.str.strip()
-    if ignore_case:
-        if col_1.dtype.kind == "O" and pd.api.types.is_string_dtype(col_1):
-            col_1 = col_1.str.upper()
-        if col_2.dtype.kind == "O" and pd.api.types.is_string_dtype(col_2):
-            col_2 = col_2.str.upper()
+
+    col_1 = normalize_string_column(
+        col_1, ignore_spaces=ignore_spaces, ignore_case=ignore_case
+    )
+    col_2 = normalize_string_column(
+        col_2, ignore_spaces=ignore_spaces, ignore_case=ignore_case
+    )
 
     # short circuit if comparing mixed type columns. Check list/arrrays or just return false for everything else.
     if pd.api.types.infer_dtype(col_1).startswith("mixed") or pd.api.types.infer_dtype(
@@ -1054,3 +1048,35 @@ def generate_id_within_group(
         )
     else:
         return dataframe[join_columns].groupby(join_columns).cumcount()
+
+
+def normalize_string_column(
+    column: pd.Series, ignore_spaces: bool, ignore_case: bool
+) -> pd.Series:
+    """Normalize a string column by converting to upper case and stripping whitespace.
+
+    Parameters
+    ----------
+    column : pd.Series
+        The column to normalize
+    ignore_spaces : bool
+        Whether to ignore spaces when normalizing
+    ignore_case : bool
+        Whether to ignore case when normalizing
+
+    Returns
+    -------
+    pd.Series
+        The normalized column
+
+    Notes
+    -----
+    Will not operate on categorical columns.
+    """
+    if (column.dtype.kind == "O" and pd.api.types.infer_dtype(column) == "string") or (
+        pd.api.types.is_string_dtype(column)
+        and not isinstance(column.dtype, pd.CategoricalDtype)
+    ):
+        column = column.str.strip() if ignore_spaces else column
+        column = column.str.upper() if ignore_case else column
+    return column
diff --git a/datacompy/polars.py b/datacompy/polars.py
index a00fbb2..12760ee 100644
--- a/datacompy/polars.py
+++ b/datacompy/polars.py
@@ -66,9 +66,9 @@ class PolarsCompare(BaseCompare):
         A string name for the second dataframe
     ignore_spaces : bool, optional
         Flag to strip whitespace (including newlines) from string columns (including any join
-        columns)
+        columns). Excludes categoricals.
     ignore_case : bool, optional
-        Flag to ignore the case of string columns
+        Flag to ignore the case of string columns. Excludes categoricals.
     cast_column_names_lower: bool, optional
         Boolean indicator that controls of column names will be cast into lower case
 
@@ -233,10 +233,19 @@ class PolarsCompare(BaseCompare):
         return OrderedSet(self.df1.columns) & OrderedSet(self.df2.columns)
 
     def _dataframe_merge(self, ignore_spaces: bool) -> None:
-        """Merge df1 to df2 on the join columns.
+        """Perform an outer join between two dataframes and categorize rows into unique and intersecting groups based on the join columns.
 
-        To get df1 - df2, df2 - df1
-        and df1 & df2.
+        Parameters
+        ----------
+        ignore_spaces : bool
+            If True, normalizes string columns by ignoring spaces during the join operation.
+
+        Notes
+        -----
+        - Temporary columns may be added to the dataframes during processing
+          and are cleaned up before final output.
+        - The method assumes that `self.df1`, `self.df2`, and `self.join_columns`
+          are properly initialized before calling this method.
         """
         params: Dict[str, Any]
         LOG.debug("Outer joining")
@@ -261,10 +270,16 @@ class PolarsCompare(BaseCompare):
 
         if ignore_spaces:
             for column in self.join_columns:
-                if str(df1[column].dtype) in STRING_TYPE:
-                    df1 = df1.with_columns(pl.col(column).str.strip_chars())
-                if str(df2[column].dtype) in STRING_TYPE:
-                    df2 = df2.with_columns(pl.col(column).str.strip_chars())
+                df1 = df1.with_columns(
+                    normalize_string_column(
+                        df1[column], ignore_spaces=ignore_spaces, ignore_case=False
+                    )
+                )
+                df2 = df2.with_columns(
+                    normalize_string_column(
+                        df2[column], ignore_spaces=ignore_spaces, ignore_case=False
+                    )
+                )
 
         df1_non_join_columns = OrderedSet(df1.columns) - OrderedSet(temp_join_columns)
         df2_non_join_columns = OrderedSet(df2.columns) - OrderedSet(temp_join_columns)
@@ -841,17 +856,8 @@ def columns_equal(
     """
     compare: pl.Series
 
-    if ignore_spaces:
-        if str(col_1.dtype) in STRING_TYPE:
-            col_1 = col_1.str.strip_chars()
-        if str(col_2.dtype) in STRING_TYPE:
-            col_2 = col_2.str.strip_chars()
-
-    if ignore_case:
-        if str(col_1.dtype) in STRING_TYPE:
-            col_1 = col_1.str.to_uppercase()
-        if str(col_2.dtype) in STRING_TYPE:
-            col_2 = col_2.str.to_uppercase()
+    col_1 = normalize_string_column(col_1, ignore_spaces, ignore_case)
+    col_2 = normalize_string_column(col_2, ignore_spaces, ignore_case)
 
     if (
         (str(col_1.dtype) in STRING_TYPE and str(col_2.dtype) in STRING_TYPE)
@@ -895,8 +901,7 @@ def columns_equal(
 def compare_string_and_date_columns(col_1: pl.Series, col_2: pl.Series) -> pl.Series:
     """Compare a string column and date column, value-wise.
 
-    This tries to
-    convert a string column to a date column and compare that way.
+    This tries to convert a string column to a date column and compare that way.
 
     Parameters
     ----------
@@ -1021,3 +1026,34 @@ def generate_id_within_group(
         return dataframe.select(
             rn=pl.col(dataframe.columns[0]).cum_count().over(join_columns)
         ).to_series()
+
+
+def normalize_string_column(
+    column: pl.Series, ignore_spaces: bool, ignore_case: bool
+) -> pl.Series:
+    """Normalize a string column by converting to upper case and stripping whitespace.
+
+    Parameters
+    ----------
+    column : pl.Series
+        The column to normalize
+    ignore_spaces : bool
+        Whether to ignore spaces when normalizing
+    ignore_case : bool
+        Whether to ignore case when normalizing
+
+    Returns
+    -------
+    pl.Series
+        The normalized column
+
+    Notes
+    -----
+    Will not operate on categorical columns.
+    """
+    if str(column.dtype.base_type()) in STRING_TYPE:
+        if ignore_spaces:
+            column = column.str.strip_chars()
+        if ignore_case:
+            column = column.str.to_uppercase()
+    return column


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index 4040c47..45c817b 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -443,16 +443,22 @@ def test_categorical_column():
             "idx": [1, 2, 3],
             "foo": ["A", "B", np.nan],
             "bar": ["A", "B", np.nan],
+            "foo_bad": ["    A   ", "B", np.nan],
         }
     )
-    for col in ("foo", "bar"):
-        df[col] = df[col].astype("category")
+    for col in ("foo", "bar", "foo_bad"):
         df[col] = df[col].astype("category")
 
     actual_out = datacompy.columns_equal(
         df.foo, df.bar, ignore_spaces=True, ignore_case=True
     )
     assert actual_out.all()
+
+    actual_out = datacompy.columns_equal(
+        df.foo, df.foo_bad, ignore_spaces=True, ignore_case=True
+    )
+    assert list(actual_out) == [False, True, True]
+
     compare = datacompy.Compare(df, df, join_columns=["idx"])
     assert compare.intersect_rows["foo_match"].all()
     assert compare.intersect_rows["bar_match"].all()
@@ -1679,3 +1685,151 @@ def test_columns_equal_lists():
     df2 = pd.DataFrame({"array_col": [[] for _ in range(10)]})
     actual = datacompy.columns_equal(df1.array_col, df2.array_col)
     assert actual.all()
+
+
+@pytest.mark.parametrize(
+    "data, ignore_spaces, ignore_case, expected",
+    [
+        # test case for categoricals, should pass through
+        # as we don't want pandas to cast categoricals into strings.
+        (
+            pd.Series(["  cat  ", "dog", "  mouse  ", None], dtype="category"),
+            True,
+            True,
+            pd.Series(["  cat  ", "dog", "  mouse  ", None], dtype="category"),
+        ),
+        # test case for mixed types
+        (
+            pd.Series(["1   ", 2.5, None, True]),
+            True,
+            True,
+            pd.Series(["1   ", 2.5, None, True]),
+        ),
+        # test case for integers
+        (pd.Series([1, 2, 3, 4]), True, True, pd.Series([1, 2, 3, 4])),
+        (pd.Series([1, 2, 3, 4]), True, False, pd.Series([1, 2, 3, 4])),
+        (pd.Series([1, 2, 3, 4]), False, True, pd.Series([1, 2, 3, 4])),
+        (pd.Series([1, 2, 3, 4]), False, False, pd.Series([1, 2, 3, 4])),
+        # test case for floats
+        (pd.Series([1.1, 2.2, 3.3, 4.4]), True, True, pd.Series([1.1, 2.2, 3.3, 4.4])),
+        (pd.Series([1.1, 2.2, 3.3, 4.4]), True, False, pd.Series([1.1, 2.2, 3.3, 4.4])),
+        (pd.Series([1.1, 2.2, 3.3, 4.4]), False, True, pd.Series([1.1, 2.2, 3.3, 4.4])),
+        (
+            pd.Series([1.1, 2.2, 3.3, 4.4]),
+            False,
+            False,
+            pd.Series([1.1, 2.2, 3.3, 4.4]),
+        ),
+        # test case for list of strings should just passthrough
+        (
+            pd.Series([["  hello  ", "WORLD", "  Foo  ", None]]),
+            True,
+            True,
+            pd.Series([["  hello  ", "WORLD", "  Foo  ", None]]),
+        ),
+        # test case for strings of object types
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="object"),
+            True,
+            True,
+            pd.Series(["HELLO", "WORLD", "FOO", None], dtype="object"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="object"),
+            True,
+            False,
+            pd.Series(["hello", "WORLD", "Foo", None], dtype="object"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="object"),
+            False,
+            True,
+            pd.Series(["  HELLO  ", "WORLD", "  FOO  ", None], dtype="object"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="object"),
+            False,
+            False,
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="object"),
+        ),
+        (
+            pd.Series(["👋", "🌍", "🍕", None], dtype="object"),
+            True,
+            True,
+            pd.Series(["👋", "🌍", "🍕", None], dtype="object"),
+        ),
+        (
+            pd.Series(["  👋  ", "🌍", "  🍕  ", None], dtype="object"),
+            False,
+            True,
+            pd.Series(["  👋  ", "🌍", "  🍕  ", None], dtype="object"),
+        ),
+        # tests for type string[python]
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[python]"),
+            True,
+            True,
+            pd.Series(["HELLO", "WORLD", "FOO", None], dtype="string[python]"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[python]"),
+            True,
+            False,
+            pd.Series(["hello", "WORLD", "Foo", None], dtype="string[python]"),
+        ),
+        (
+            pd.Series(["👋", "🌍", "🍕", None], dtype="string[python]"),
+            True,
+            True,
+            pd.Series(["👋", "🌍", "🍕", None], dtype="string[python]"),
+        ),
+        (
+            pd.Series(["  👋  ", "🌍", "  🍕  ", None], dtype="string[python]"),
+            False,
+            True,
+            pd.Series(["  👋  ", "🌍", "  🍕  ", None], dtype="string[python]"),
+        ),
+        # tests for type string[pyarrow]
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[pyarrow]"),
+            True,
+            True,
+            pd.Series(["HELLO", "WORLD", "FOO", None], dtype="string[pyarrow]"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[pyarrow]"),
+            True,
+            False,
+            pd.Series(["hello", "WORLD", "Foo", None], dtype="string[pyarrow]"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[pyarrow]"),
+            False,
+            True,
+            pd.Series(["  HELLO  ", "WORLD", "  FOO  ", None], dtype="string[pyarrow]"),
+        ),
+        (
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[pyarrow]"),
+            False,
+            False,
+            pd.Series(["  hello  ", "WORLD", "  Foo  ", None], dtype="string[pyarrow]"),
+        ),
+        (
+            pd.Series(["👋", "🌍", "🍕", None], dtype="string[pyarrow]"),
+            True,
+            True,
+            pd.Series(["👋", "🌍", "🍕", None], dtype="string[pyarrow]"),
+        ),
+        (
+            pd.Series(["  👋  ", "🌍", "  🍕  ", None], dtype="string[pyarrow]"),
+            False,
+            True,
+            pd.Series(["  👋  ", "🌍", "  🍕  ", None], dtype="string[pyarrow]"),
+        ),
+    ],
+)
+def test_normalize_string_column(data, ignore_spaces, ignore_case, expected):
+    result = datacompy.core.normalize_string_column(
+        data, ignore_spaces=ignore_spaces, ignore_case=ignore_case
+    )
+    assert_series_equal(result, expected, check_names=False)
diff --git a/tests/test_polars.py b/tests/test_polars.py
index d4c2e77..8d3eb7a 100644
--- a/tests/test_polars.py
+++ b/tests/test_polars.py
@@ -36,6 +36,7 @@ from datacompy.polars import (
     calculate_max_diff,
     columns_equal,
     generate_id_within_group,
+    normalize_string_column,
     temp_column_name,
 )
 from polars.exceptions import ComputeError, DuplicateError, SchemaError
@@ -1450,15 +1451,27 @@ def test_categorical_column():
             "idx": [1, 2, 3],
             "foo": ["A", "B", np.nan],
             "bar": ["A", "B", np.nan],
+            "foo_bad": ["    A   ", "B", np.nan],
         },
         strict=False,
-        schema={"idx": pl.Int32, "foo": pl.Categorical, "bar": pl.Categorical},
+        schema={
+            "idx": pl.Int32,
+            "foo": pl.Categorical,
+            "bar": pl.Categorical,
+            "foo_bad": pl.Categorical,
+        },
     )
 
     actual_out = columns_equal(
         df["foo"], df["bar"], ignore_spaces=True, ignore_case=True
     )
     assert actual_out.all()
+
+    actual_out = columns_equal(
+        df["foo"], df["foo_bad"], ignore_spaces=True, ignore_case=True
+    )
+    assert list(actual_out) == [False, True, True]
+
     compare = PolarsCompare(df, df, join_columns=["idx"])
     assert compare.intersect_rows["foo_match"].all()
     assert compare.intersect_rows["bar_match"].all()
@@ -1648,3 +1661,95 @@ def test_columns_equal_lists():
     )
     actual = columns_equal(df1["array_col"], df2["array_col"])
     assert actual.all()
+
+
+@pytest.mark.parametrize(
+    "input_data, ignore_spaces, ignore_case, expected",
+    [  # Categorical datatype should just passthough
+        (
+            pl.Series(["  cat  ", "dog", "  mouse  ", None], dtype=pl.Categorical),
+            True,
+            True,
+            pl.Series(["  cat  ", "dog", "  mouse  ", None], dtype=pl.Categorical),
+        ),
+        # mixed types
+        (
+            pl.Series(["  hello  ", 1, 2.5, None], strict=False),
+            True,
+            True,
+            pl.Series(["HELLO", 1, 2.5, None], strict=False),
+        ),
+        # test case for integers
+        (pl.Series([1, 2, 3, 4]), True, True, pl.Series([1, 2, 3, 4])),
+        (pl.Series([1, 2, 3, 4]), True, False, pl.Series([1, 2, 3, 4])),
+        (pl.Series([1, 2, 3, 4]), False, True, pl.Series([1, 2, 3, 4])),
+        (pl.Series([1, 2, 3, 4]), False, False, pl.Series([1, 2, 3, 4])),
+        # test case for floats
+        (pl.Series([1.1, 2.2, 3.3, 4.4]), True, True, pl.Series([1.1, 2.2, 3.3, 4.4])),
+        (pl.Series([1.1, 2.2, 3.3, 4.4]), True, False, pl.Series([1.1, 2.2, 3.3, 4.4])),
+        (pl.Series([1.1, 2.2, 3.3, 4.4]), False, True, pl.Series([1.1, 2.2, 3.3, 4.4])),
+        (
+            pl.Series([1.1, 2.2, 3.3, 4.4]),
+            False,
+            False,
+            pl.Series([1.1, 2.2, 3.3, 4.4]),
+        ),
+        # list of strings should just passthrough
+        (
+            pl.Series([["  hello  ", "WORLD", "  Foo  ", None]]),
+            True,
+            True,
+            pl.Series([["  hello  ", "WORLD", "  Foo  ", None]]),
+        ),
+        # array of strings should just passthrough
+        (
+            pl.Series([["  hello  ", "WORLD", "  Foo  ", None]]),
+            True,
+            True,
+            pl.Series([["  hello  ", "WORLD", "  Foo  ", None]]),
+        ),
+        # strings
+        (
+            pl.Series(["  hello  ", "WORLD", "  Foo  ", None]),
+            True,
+            True,
+            pl.Series(["HELLO", "WORLD", "FOO", None]),
+        ),
+        (
+            pl.Series(["  hello  ", "WORLD", "  Foo  ", None]),
+            True,
+            False,
+            pl.Series(["hello", "WORLD", "Foo", None]),
+        ),
+        (
+            pl.Series(["  hello  ", "WORLD", "  Foo  ", None]),
+            False,
+            True,
+            pl.Series(["  HELLO  ", "WORLD", "  FOO  ", None]),
+        ),
+        (
+            pl.Series(["  hello  ", "WORLD", "  Foo  ", None]),
+            False,
+            False,
+            pl.Series(["  hello  ", "WORLD", "  Foo  ", None]),
+        ),
+        # emoji
+        (
+            pl.Series(["👋", "🌍", "🍕", None]),
+            True,
+            True,
+            pl.Series(["👋", "🌍", "🍕", None]),
+        ),
+        (
+            pl.Series(["  👋  ", "🌍", "  🍕  ", None]),
+            False,
+            True,
+            pl.Series(["  👋  ", "🌍", "  🍕  ", None]),
+        ),
+    ],
+)
+def test_normalize_string_column(input_data, ignore_spaces, ignore_case, expected):
+    result = normalize_string_column(
+        input_data, ignore_spaces=ignore_spaces, ignore_case=ignore_case
+    )
+    assert_series_equal(result, expected, check_names=False)