diff --git a/diffly/_conditions.py b/diffly/_conditions.py index 98e99b0..a86cd97 100644 --- a/diffly/_conditions.py +++ b/diffly/_conditions.py @@ -153,8 +153,10 @@ def _compare_columns( ) return col_left.eq_missing(col_right) - if _different_enums(dtype_left, dtype_right) or _enum_and_categorical( - dtype_left, dtype_right + if ( + _different_enums(dtype_left, dtype_right) + or _different_categoricals(dtype_left, dtype_right) + or _enum_and_categorical(dtype_left, dtype_right) ): # Enums with different categories as well as enums and categoricals # can't be compared directly. @@ -267,6 +269,7 @@ def _needs_element_wise_comparison( _is_float_numeric_pair(dtype_left, dtype_right) or _is_temporal_pair(dtype_left, dtype_right) or _different_enums(dtype_left, dtype_right) + or _different_categoricals(dtype_left, dtype_right) or _enum_and_categorical(dtype_left, dtype_right) ): return True @@ -322,6 +325,16 @@ def _different_enums( return isinstance(left, pl.Enum) and isinstance(right, pl.Enum) and left != right +def _different_categoricals( + left: DataType | DataTypeClass, right: DataType | DataTypeClass +) -> bool: + return ( + isinstance(left, pl.Categorical) + and isinstance(right, pl.Categorical) + and left != right + ) + + def _enum_and_categorical( left: DataType | DataTypeClass, right: DataType | DataTypeClass ) -> bool: diff --git a/tests/test_conditions.py b/tests/test_conditions.py index 5895be8..1f34c49 100644 --- a/tests/test_conditions.py +++ b/tests/test_conditions.py @@ -555,6 +555,43 @@ def test_condition_equal_columns_list_of_different_enums() -> None: assert actual.to_list() == [True, False] +def test_condition_equal_columns_different_categorical() -> None: + # Arrange + fruits = pl.Categorical(categories=pl.Categories(name="fruits")) + fruit = pl.Categorical(categories=pl.Categories(name="fruit")) + + lhs = pl.DataFrame( + {"pk": [1, 2], "a": ["apple", "orange"]}, + schema_overrides={"a": fruits}, + ) + rhs = pl.DataFrame( + {"pk": [1, 2], "a": ["apple", "banana"]}, + schema_overrides={"a": fruit}, + ) + c = compare_frames(lhs, rhs, primary_key="pk") + + # Act + lhs = lhs.rename({"a": "a_left"}) + rhs = rhs.rename({"a": "a_right"}) + actual = ( + lhs.join(rhs, on="pk", maintain_order="left") + .select( + condition_equal_columns( + "a", + dtype_left=lhs.schema["a_left"], + dtype_right=rhs.schema["a_right"], + max_list_length=None, + abs_tol=c.abs_tol_by_column["a"], + rel_tol=c.rel_tol_by_column["a"], + ) + ) + .to_series() + ) + + # Assert + assert actual.to_list() == [True, False] + + @pytest.mark.parametrize( ("dtype_left", "dtype_right", "can_compare_dtypes"), [