Repository: unionai-oss/pandera

Issue:
Custom check fails with `pl.DataFrame`
- [x] I have checked that this issue has not already been reported.
- [x] I have confirmed this bug exists on the latest version of pandera.
- [ ] (optional) I have confirmed this bug exists on the master branch of pandera.

#### Code Sample, a copy-pastable example

I'm seeing an error when trying to perform schema validation with custom checks on a polars dataframe: `pandera.errors.SchemaError: ComputeError("custom python function failed: cannot unpack series of type 'str' into 'bool'")`

```python
import polars as pl
import pandera.polars as pa


# Custom check function
def check_len(v: str) -> bool:
    return len(v) <= 20

schema = pa.DataFrameSchema(
    {
        "phone": pa.Column(
            dtype=str,
            checks=pa.Check.str_matches(r"\(\d{3}\) \d{3}-\d{4}"),
        ),
        "fruit": pa.Column(
            dtype=str,
            checks=pa.Check(check_len, element_wise=True),
        ),
    }
)

df = pl.DataFrame(
    {
        "phone": ["(123) 234-2342", "(213) 234-2345", "(213) 234-2345"],
        "fruit": ["apple", "pear", "banana"],
    }
)

df.pipe(schema.validate)
# C:\local\.venv\Lib\site-packages\pandera\backends\polars\base.py:74: MapWithoutReturnDtypeWarning: Calling `map_elements` without specifying `return_dtype` can lead to unpredictable results. Specify `return_dtype` to silence this warning.
#   passed = check_result.check_passed.collect().item()
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
#   File "C:\local\.venv\Lib\site-packages\polars\dataframe\frame.py", line 5150, in pipe
#     return function(self, *args, **kwargs)
#            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#   File "C:\local\.venv\Lib\site-packages\pandera\api\polars\container.py", line 58, in validate
#     output = self.get_backend(check_obj).validate(
#              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#   File "C:\local\.venv\Lib\site-packages\pandera\backends\polars\container.py", line 114, in validate
#     error_handler.collect_error(
#   File "C:\local\.venv\Lib\site-packages\pandera\api\base\error_handler.py", line 54, in collect_error
#     raise schema_error from original_exc
#   File "C:\local\.venv\Lib\site-packages\pandera\backends\polars\container.py", line 182, in run_schema_component_checks
#     result = schema_component.validate(check_obj, lazy=lazy)
#              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#   File "C:\local\.venv\Lib\site-packages\pandera\api\polars\components.py", line 141, in validate
#     output = self.get_backend(check_obj).validate(
#              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#   File "C:\local\.venv\Lib\site-packages\pandera\backends\polars\components.py", line 81, in validate
#     error_handler = self.run_checks_and_handle_errors(
#                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#   File "C:\local\.venv\Lib\site-packages\pandera\backends\polars\components.py", line 147, in run_checks_and_handle_errors
#     error_handler.collect_error(
#   File "C:\local\.venv\Lib\site-packages\pandera\api\base\error_handler.py", line 54, in collect_error
#     raise schema_error from original_exc
# pandera.errors.SchemaError: ComputeError("custom python function failed: cannot unpack series of type `str` into `bool`")```
```

Removing the phone check will avoid this error, but then a warning about a missing `return_dtype` value is returned:

```python
import polars as pl
import pandera.polars as pa

# Custom check function
def check_len(v: str) -> bool:
    return len(v) <= 20

schema = pa.DataFrameSchema(
    {
        "fruit": pa.Column(
            dtype=str,
            checks=pa.Check(check_len, element_wise=True),
        ),
    }
)

df = pl.DataFrame(
    {
        "fruit": ["apple", "pear", "banana"],
    }
)

df.pipe(schema.validate)
# C:\local\.venv\Lib\site-packages\pandera\backends\polars\base.py:74: MapWithoutReturnDtypeWarning: Calling `map_elements` without specifying `return_dtype` can lead to unpredictable results. Specify `return_dtype` to silence this warning.
#   passed = check_result.check_passed.collect().item()
# C:\local\.venv\Lib\site-packages\pandera\backends\polars\base.py:112: MapWithoutReturnDtypeWarning: Calling `map_elements` without specifying `return_dtype` can lead to unpredictable results. Specify `return_dtype` to silence this warning.
#   check_output=check_result.check_output.collect(),
# shape: (3, 1)
# ┌────────┐
# │ fruit  │
# │ ---    │
# │ str    │
# ╞════════╡
# │ apple  │
# │ pear   │
# │ banana │
# └────────┘
```

#### Expected behavior
Schema validation should run successfully for this DataFrame when custom checks are included alongside built-in check methods. 

#### Desktop (please complete the following information):

 - OS: Windows 10
 - Browser: Chrome
 - Version `pandera==0.19.0b1`


Discussion:
~It's difficult to trust that custom checks are actually working with the `LazyFrame` in this reprex after discovering #1566.~

Edit: [LazyFrame validation does not apply to data-level checks by design.](https://github.com/unionai-oss/pandera/issues/1566#issuecomment-2050719906) 

Resolving patch:
diff --git a/pandera/__init__.py b/pandera/__init__.py
--- a/pandera/__init__.py
+++ b/pandera/__init__.py
@@ -69,6 +69,7 @@
 import pandera.backends.base.builtin_checks
 import pandera.backends.base.builtin_hypotheses
 import pandera.backends.pandas
+from pandera.backends.pandas.register import register_pandas_backends
 
 from pandera.schema_inference.pandas import infer_schema
 from pandera.version import __version__
@@ -178,3 +179,6 @@
     # version
     "__version__",
 ]
+
+
+register_pandas_backends()
diff --git a/pandera/backends/polars/checks.py b/pandera/backends/polars/checks.py
--- a/pandera/backends/polars/checks.py
+++ b/pandera/backends/polars/checks.py
@@ -44,9 +44,10 @@ def preprocess(self, check_obj: pl.LazyFrame, key: Optional[str]):
     def apply(self, check_obj: PolarsData):
         """Apply the check function to a check object."""
         if self.check.element_wise:
+            selector = pl.col(check_obj.key or "*")
             out = check_obj.lazyframe.with_columns(
-                pl.col(check_obj.key or "*").map_elements(self.check_fn)
-            )
+                selector.map_elements(self.check_fn, return_dtype=pl.Boolean)
+            ).select(selector)
         else:
             out = self.check_fn(check_obj)
 
diff --git a/pandera/polars.py b/pandera/polars.py
--- a/pandera/polars.py
+++ b/pandera/polars.py
@@ -11,6 +11,8 @@
 from pandera.api.polars.container import DataFrameSchema
 from pandera.api.polars.model import DataFrameModel
 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
 
-import pandera.backends.polars
+
+register_polars_backends()


Test patch:
diff --git a/tests/polars/test_polars_check.py b/tests/polars/test_polars_check.py
--- a/tests/polars/test_polars_check.py
+++ b/tests/polars/test_polars_check.py
@@ -151,3 +151,33 @@ def test_polars_element_wise_dataframe_check(lf):
             schema.validate(invalid_lf)
         except pa.errors.SchemaError as exc:
             exc.failure_cases.equals(pl.DataFrame({col: [-1, -4]}))
+
+
+def test_polars_element_wise_dataframe_different_dtypes(column_lf):
+
+    # Custom check function
+    def check_gt_2(v: int) -> bool:
+        return v > 2
+
+    def check_len_ge_2(v: str) -> bool:
+        return len(v) >= 2
+
+    lf = column_lf.with_columns(
+        str_col=pl.Series(["aaa", "bb", "c", "dd"], dtype=str)
+    )
+
+    schema = pa.DataFrameSchema(
+        {
+            "col": pa.Column(
+                dtype=int, checks=pa.Check(check_gt_2, element_wise=True)
+            ),
+            "str_col": pa.Column(
+                dtype=str, checks=pa.Check(check_len_ge_2, element_wise=True)
+            ),
+        }
+    )
+
+    try:
+        schema.validate(lf, lazy=True)
+    except pa.errors.SchemaErrors as exc:
+        assert exc.failure_cases["failure_case"].to_list() == ["1", "2", "c"]