diff --git a/bulwark/checks.py b/bulwark/checks.py index aadaf5b..4bade7b 100644 --- a/bulwark/checks.py +++ b/bulwark/checks.py @@ -398,3 +398,28 @@ def custom_check(check_func, df, *args, **kwargs): raise return df + + +def column_compare(df, col_a, col_b, func): + """Asserts that two columns adhere to the condition. + + Args: + df (pd.DataFrame): Any pd.DataFrame. + col_a (str): First column name. + col_b (str): Second column name. + func (): Function to compare columns. + + Returns: + Original `df`. + + """ + try: + vfunc = np.vectorize(func) + result = vfunc(df[col_a], df[col_b]) + + if not np.all(result): + raise AssertionError("Column comparison failed. Failed column indices: {}".format(np.argwhere(result!=True))) + except Exception as e: + raise + + return df diff --git a/tests/test_checks.py b/tests/test_checks.py index e60389f..8fc994c 100644 --- a/tests/test_checks.py +++ b/tests/test_checks.py @@ -390,3 +390,11 @@ def f(df, length): with pytest.raises(AssertionError): dc.CustomCheck(f, 4)(_noop)(df) + + +def test_column_compare(): + df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + + tm.assert_frame_equal(df, ck.column_compare(df, 'b', 'a', lambda x, y: x > y)) + with pytest.raises(AssertionError): + result = ck.column_compare(df, 'a', 'b', lambda x, y: x > y)