Support table-level checks, including for built-in

Signed-off-by: Deepyaman Datta <[email protected]>

Author: deepyaman

pandera/backends/ibis/builtin_checks.py

@@ -11,8 +11,6 @@
 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")
 
@@ -46,9 +44,7 @@ def equal_to(data: IbisData, value: Any) -> ir.Table:
         equal to this value.
     """
     value = _infer_interval_with_mixed_units(value)
-    return data.table.mutate(
-        s.across(_selector(data.key), _ == value, names=check_col_name)
-    )
+    return data.table.select(s.across(_selector(data.key), _ == value))
 
 
 @register_builtin_check(
@@ -63,9 +59,7 @@ 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.mutate(
-        s.across(_selector(data.key), _ != value, names=check_col_name)
-    )
+    return data.table.select(s.across(_selector(data.key), _ != value))
 
 
 @register_builtin_check(
@@ -82,9 +76,7 @@ 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.mutate(
-        s.across(_selector(data.key), _ > value, names=check_col_name)
-    )
+    return data.table.select(s.across(_selector(data.key), _ > value))
 
 
 @register_builtin_check(
@@ -100,9 +92,7 @@ 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.mutate(
-        s.across(_selector(data.key), _ >= value, names=check_col_name)
-    )
+    return data.table.select(s.across(_selector(data.key), _ >= value))
 
 
 @register_builtin_check(
@@ -119,9 +109,7 @@ 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.mutate(
-        s.across(_selector(data.key), _ < value, names=check_col_name)
-    )
+    return data.table.select(s.across(_selector(data.key), _ < value))
 
 
 @register_builtin_check(
@@ -137,6 +125,4 @@ 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.mutate(
-        s.across(_selector(data.key), _ <= value, names=check_col_name)
-    )
+    return data.table.select(s.across(_selector(data.key), _ <= value))

pandera/backends/ibis/checks.py

@@ -65,6 +65,24 @@ def apply(self, check_obj: IbisData):
                 out = check_obj.table.mutate(
                     **{f"{k}{CHECK_OUTPUT_SUFFIX}": v for k, v in out.items()}
                 )
+            elif isinstance(out, ir.Table):
+                out = out.rename(f"{{name}}{CHECK_OUTPUT_SUFFIX}")
+                try:
+                    out = check_obj.table.join(out, how="positional")
+                except Exception:  # Backend doesn't support positional join
+                    index_col = "__idx__"
+                    out = (
+                        check_obj.table.mutate(
+                            **{index_col: ibis.row_number().over()}
+                        )
+                        .join(
+                            out.mutate(
+                                **{index_col: ibis.row_number().over()}
+                            ),
+                            index_col,
+                        )
+                        .drop(index_col)
+                    )
 
         if isinstance(out, ir.Table):
             # for checks that return a boolean table, make sure all columns

pandera/constants.py

@@ -3,4 +3,3 @@
 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}"

tests/ibis/test_ibis_check.py

@@ -7,7 +7,10 @@
 import pandas as pd
 import ibis
 import ibis.expr.types as ir
+import ibis.selectors as s
 import pandera.ibis as pa
+from ibis import _, selectors as s
+
 from pandera.backends.ibis.register import register_ibis_backends
 from pandera.constants import CHECK_OUTPUT_KEY
 
@@ -66,9 +69,7 @@ def test_ibis_column_check(
 
 
 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}
-    )
+    return data.table.select(s.across(s.numeric(), _ >= 0))
 
 
 def _df_check_fn_dict_out(data: pa.IbisData) -> Dict[str, ir.BooleanColumn]: