built-in checks support table-level checks

Signed-off-by: cosmicBboy <[email protected]>

Author: cosmicBboy

pandera/backends/ibis/builtin_checks.py

@@ -1,13 +1,18 @@
 """Built-in checks for Ibis."""
 
 import datetime
-from typing import Any, TypeVar
+from typing import Any, Optional, TypeVar
 
 import ibis
 import ibis.expr.types as ir
+from ibis import _, selectors as s
+from ibis.common.selectors import Selector
 
 from pandera.api.extensions import register_builtin_check
 from pandera.api.ibis.types import IbisData
+from pandera.backends.ibis.utils import select_column
+from pandera.constants import check_col_name
+
 
 T = TypeVar("T")
 
@@ -24,6 +29,10 @@ def _infer_interval_with_mixed_units(value: Any) -> Any:
     return value
 
 
+def _selector(key: Optional[str]) -> Selector:
+    return s.all() if key is None else select_column(key)
+
+
 @register_builtin_check(
     aliases=["eq"],
     error="equal_to({value})",
@@ -37,7 +46,9 @@ def equal_to(data: IbisData, value: Any) -> ir.Table:
         equal to this value.
     """
     value = _infer_interval_with_mixed_units(value)
-    return data.table[data.key] == value
+    return data.table.mutate(
+        s.across(_selector(data.key), _ == value, names=check_col_name)
+    )
 
 
 @register_builtin_check(
@@ -52,12 +63,14 @@ def not_equal_to(data: IbisData, value: Any) -> ir.Table:
     :param value: This value must not occur in the checked data structure.
     """
     value = _infer_interval_with_mixed_units(value)
-    return data.table[data.key] != value
+    return data.table.mutate(
+        s.across(_selector(data.key), _ != value, names=check_col_name)
+    )
 
 
 @register_builtin_check(
     aliases=["gt"],
-    error="greater_than({value})",
+    error="greater_than({min_value})",
 )
 def greater_than(data: IbisData, min_value: Any) -> ir.Table:
     """Ensure values of a column are strictly greater than a minimum
@@ -69,12 +82,14 @@ def greater_than(data: IbisData, min_value: Any) -> ir.Table:
         to the dtype of the :class:`ir.Column` to be validated.
     """
     value = _infer_interval_with_mixed_units(min_value)
-    return data.table[data.key] > value
+    return data.table.mutate(
+        s.across(_selector(data.key), _ > value, names=check_col_name)
+    )
 
 
 @register_builtin_check(
     aliases=["ge"],
-    error="greater_than_or_equal_to({value})",
+    error="greater_than_or_equal_to({min_value})",
 )
 def greater_than_or_equal_to(data: IbisData, min_value: Any) -> ir.Table:
     """Ensure all values are greater than or equal to a minimum value.
@@ -85,7 +100,9 @@ def greater_than_or_equal_to(data: IbisData, min_value: Any) -> ir.Table:
         to the dtype of the :class:`ir.Column` to be validated.
     """
     value = _infer_interval_with_mixed_units(min_value)
-    return data.table[data.key] >= value
+    return data.table.mutate(
+        s.across(_selector(data.key), _ >= value, names=check_col_name)
+    )
 
 
 @register_builtin_check(
@@ -102,7 +119,9 @@ def less_than(data: IbisData, max_value: Any) -> ir.Table:
         :class:`ir.Column` to be validated.
     """
     value = _infer_interval_with_mixed_units(max_value)
-    return data.table[data.key] < value
+    return data.table.mutate(
+        s.across(_selector(data.key), _ < value, names=check_col_name)
+    )
 
 
 @register_builtin_check(
@@ -118,4 +137,6 @@ def less_than_or_equal_to(data: IbisData, max_value: Any) -> ir.Table:
         :class:`ir.Column` to be validated.
     """
     value = _infer_interval_with_mixed_units(max_value)
-    return data.table[data.key] <= value
+    return data.table.mutate(
+        s.across(_selector(data.key), _ <= value, names=check_col_name)
+    )

pandera/backends/ibis/checks.py

@@ -13,9 +13,7 @@
 from pandera.api.ibis.types import IbisData
 from pandera.backends.base import BaseCheckBackend
 from pandera.backends.ibis.utils import select_column
-from pandera.constants import CHECK_OUTPUT_KEY
-
-CHECK_OUTPUT_SUFFIX = f"__{CHECK_OUTPUT_KEY}__"
+from pandera.constants import CHECK_OUTPUT_KEY, CHECK_OUTPUT_SUFFIX
 
 
 class IbisCheckBackend(BaseCheckBackend):
@@ -69,15 +67,15 @@ def apply(self, check_obj: IbisData):
                 )
 
         if isinstance(out, ir.Table):
-            # for checks that return a boolean dataframe, make sure all columns
+            # for checks that return a boolean table, make sure all columns
             # are boolean and reduce to a single boolean column.
             acc = ibis.literal(True)
             for col in out.columns:
                 if col.endswith(CHECK_OUTPUT_SUFFIX):
                     assert out[col].type().is_boolean(), (
                         f"column '{col[: -len(CHECK_OUTPUT_SUFFIX)]}' "
                         "is not boolean. If check function returns a "
-                        "dataframe, it must contain only boolean columns."
+                        "table, it must contain only boolean columns."
                     )
                     acc = acc & out[col]
             return out.mutate({CHECK_OUTPUT_KEY: acc})
@@ -145,7 +143,6 @@ def postprocess_table_output(
             s.endswith(f"__{CHECK_OUTPUT_KEY}__")
             | select_column(CHECK_OUTPUT_KEY)
         )
-
         if check_obj.key is not None:
             failure_cases = failure_cases.select(check_obj.key)
         return CheckResult(

pandera/constants.py

@@ -1,4 +1,6 @@
 """Pandera constants."""
 
 CHECK_OUTPUT_KEY = "check_output"
+CHECK_OUTPUT_SUFFIX = f"__{CHECK_OUTPUT_KEY}__"
 FAILURE_CASE_KEY = "failure_case"
+check_col_name = f"{{col}}{CHECK_OUTPUT_SUFFIX}"

pandera/polars.py

@@ -13,6 +13,7 @@
 from pandera.api.polars.types import PolarsData
 from pandera.backends.polars.register import register_polars_backends
 from pandera.decorators import check_input, check_io, check_output, check_types
+from pandera.constants import check_col_name
 
 register_polars_backends()
 
@@ -31,4 +32,5 @@
     "errors",
     "Field",
     "PolarsData",
+    "check_col_name",
 ]

tests/ibis/test_ibis_check.py

@@ -65,6 +65,12 @@ def test_ibis_column_check(
         assert check_output == expected_output
 
 
+def _df_check_fn_table_out(data: pa.IbisData) -> ir.Table:
+    return data.table.mutate(
+        {col: data.table[col] >= 0 for col in data.table.columns}
+    )
+
+
 def _df_check_fn_dict_out(data: pa.IbisData) -> Dict[str, ir.BooleanColumn]:
     return {col: data.table[col] >= 0 for col in data.table.columns}
 
@@ -83,6 +89,14 @@ def _df_check_fn_scalar_out(data: pa.IbisData) -> ir.BooleanScalar:
 @pytest.mark.parametrize(
     "check_fn, invalid_data, expected_output",
     [
+        [
+            _df_check_fn_table_out,
+            {
+                "col_1": [-1, 2, -3, 4],
+                "col_2": [1, 2, 3, -4],
+            },
+            [False, True, False, False],
+        ],
         [
             _df_check_fn_dict_out,
             {