bugfix: correctly support dataframe-level polars checks

pandera/api/polars/types.py

@@ -1,13 +1,13 @@
 """Polars types."""
 
-from typing import NamedTuple, Optional, Union, TypeVar
+from typing import NamedTuple, Union, TypeVar
 
 import polars as pl
 
 
 class PolarsData(NamedTuple):
     lazyframe: pl.LazyFrame
-    key: Optional[str] = None
+    key: str = "*"
 
 
 class CheckResult(NamedTuple):

pandera/backends/polars/base.py

@@ -8,7 +8,10 @@
 
 from pandera.api.base.error_handler import ErrorHandler
 from pandera.api.polars.types import CheckResult, PolarsFrame
-from pandera.api.polars.utils import get_lazyframe_column_dtypes
+from pandera.api.polars.utils import (
+    get_lazyframe_column_dtypes,
+    get_lazyframe_schema,
+)
 from pandera.backends.base import BaseSchemaBackend, CoreCheckResult
 from pandera.constants import CHECK_OUTPUT_KEY
 from pandera.errors import (
@@ -92,7 +95,11 @@ def run_check(
                 )
             else:
                 # use check_result
-                failure_cases = check_result.failure_cases.collect()
+                _failure_cases = check_result.failure_cases
+                if CHECK_OUTPUT_KEY in get_lazyframe_schema(_failure_cases):
+                    _failure_cases = _failure_cases.drop(CHECK_OUTPUT_KEY)
+
+                failure_cases = _failure_cases.collect()
                 failure_cases_msg = failure_cases.head().rows(named=True)
                 message = (
                     f"{schema.__class__.__name__} '{schema.name}' failed "

pandera/backends/polars/checks.py

@@ -132,6 +132,6 @@ def __call__(
         key: Optional[str] = None,
     ) -> CheckResult:
         check_obj = self.preprocess(check_obj, key)
-        polars_data = PolarsData(check_obj, key)
+        polars_data = PolarsData(check_obj, key or "*")
         check_output = self.apply(polars_data)
         return self.postprocess(polars_data, check_output)

tests/polars/test_polars_container.py

@@ -695,3 +695,21 @@ def test_dataframe_column_level_coerce():
 
     schema = schema.update_column("b", coerce=True)
     assert_frame_equal(schema.validate(df), df.cast({"a": int, "b": float}))
+
+
+def test_dataframe_level_check():
+    schema = DataFrameSchema(
+        {
+            "a": Column(int),
+            "b": Column(int),
+            "c": Column(int),
+        },
+        checks=[pa.Check.gt(0)],
+    )
+
+    lf = pl.LazyFrame({"a": [-1, 2, 3], "b": [4, -5, 6], "c": [7, 8, -9]})
+    try:
+        schema.validate(lf.collect(), lazy=True)
+    except pa.errors.SchemaErrors as exc:
+        # expect all rows to fail
+        assert exc.failure_cases.shape[0] == 3