column compare

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>): 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

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)