Repository: unionai-oss/pandera

Issue:
BUG: Polars DataFrameModel.validate crashes with `sample` specified
**Describe the bug**
A clear and concise description of what the bug is.

- [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.
- [x] (optional) I have confirmed this bug exists on the main branch of pandera.



#### Code Sample, a copy-pastable example

```python
import pandera.polars as pa

import pandas as pd
import polars as pl

from pandera.typing import Series


class SchemaPolars(pa.DataFrameModel):
    col1: Series[int]
    col2: Series[int]

pandas_df = pd.DataFrame({"col1": [1, 2, 3], "col2": [1, 2, 3]})
polars_df = pl.from_pandas(pandas_df)

result2 = SchemaPolars.validate(polars_df, sample=10)
Traceback (most recent call last):
  File "C:\Data\myDocuments\Code\python_other\pandera\fork\tester_sample.py", line 24, in <module>
    result2 = SchemaPolars.validate(lazyframe, sample=10)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Data\myDocuments\Code\python_other\pandera\fork\pandera\api\dataframe\model.py", line 289, in validate
    cls.to_schema().validate(
  File "C:\Data\myDocuments\Code\python_other\pandera\fork\pandera\api\polars\container.py", line 64, in validate
    output = self.get_backend(check_obj).validate(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Data\myDocuments\Code\python_other\pandera\fork\pandera\backends\polars\container.py", line 80, in validate
    sample = self.subsample(check_obj, head, tail, sample, random_state)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Data\myDocuments\Code\python_other\pandera\fork\pandera\backends\polars\base.py", line 50, in subsample
    check_obj.sample(sample, random_state=random_state)
    ^^^^^^^^^^^^^^^^
AttributeError: 'LazyFrame' object has no attribute 'sample'
```

#### Expected behavior
Sample should work, or should raise a NotImplementedError if not supported.

#### Desktop (please complete the following information):

 - OS: Windows 11, using python 3.12.8

#### Additional context
I came across this looking at running mypy over pandera.
The implementation calls .sample which is a pl.DataFrame method, but there is no lazy equivalent.

https://github.com/pola-rs/polars/issues/3933 discusses this with some potential workarounds listed, e.g.
`lazy_df.with_row_index().filter(col("index").hash(seed)%10 == 1).drop("index")`


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss) Report
Attention: Patch coverage is `0%` with `16 lines` in your changes missing coverage. Please review.
> Project coverage is 59.60%. Comparing base [(`812b2a8`)](https://app.codecov.io/gh/unionai-oss/pandera/commit/812b2a8af8c123c46caa348837da04e8c6573260?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss) to head [(`5c0f037`)](https://app.codecov.io/gh/unionai-oss/pandera/commit/5c0f03700451b8cb74210a072318a07e4a442ec2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss).
> Report is 199 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss) | Patch % | Lines |
|---|---|---|
| [pandera/backends/polars/container.py](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?src=pr&el=tree&filepath=pandera%2Fbackends%2Fpolars%2Fcontainer.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss#diff-cGFuZGVyYS9iYWNrZW5kcy9wb2xhcnMvY29udGFpbmVyLnB5) | 0.00% | [11 Missing :warning: ](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss) |
| [pandera/api/polars/container.py](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?src=pr&el=tree&filepath=pandera%2Fapi%2Fpolars%2Fcontainer.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss#diff-cGFuZGVyYS9hcGkvcG9sYXJzL2NvbnRhaW5lci5weQ==) | 0.00% | [3 Missing :warning: ](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss) |
| [pandera/backends/polars/base.py](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?src=pr&el=tree&filepath=pandera%2Fbackends%2Fpolars%2Fbase.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss#diff-cGFuZGVyYS9iYWNrZW5kcy9wb2xhcnMvYmFzZS5weQ==) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss) |

> :exclamation:  There is a different number of reports uploaded between BASE (812b2a8) and HEAD (5c0f037). Click for more details.
> 
> <details><summary>HEAD has 43 uploads less than BASE</summary>
>
>| Flag | BASE (812b2a8) | HEAD (5c0f037) |
>|------|------|------|
>||48|5|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main    #1923       +/-   ##
===========================================
- Coverage   94.28%   59.60%   -34.68%     
===========================================
  Files          91      121       +30     
  Lines        7013     9382     +2369     
===========================================
- Hits         6612     5592     -1020     
- Misses        401     3790     +3389     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/unionai-oss/pandera/pull/1923?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=unionai-oss).


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index aa35430..eee4d05 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -44,7 +44,7 @@ repos:
         args: [--py38-plus, --keep-runtime-typing]
 
   - repo: https://github.com/pycqa/pylint
-    rev: v3.2.7
+    rev: v3.3.6
     hooks:
       - id: pylint
         args: ["--disable=import-error"]
diff --git a/.pylintrc b/.pylintrc
index f1d8695..f8010a0 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -56,5 +56,6 @@ disable=
     logging-fstring-interpolation,
     multiple-statements,
     cyclic-import,
+    too-many-positional-arguments,
     # Due to custom `immutable` decorator replacing `dataclasses.dataclass`
     invalid-field-call
diff --git a/docs/source/fastapi.md b/docs/source/fastapi.md
index 365a22c..287cf52 100644
--- a/docs/source/fastapi.md
+++ b/docs/source/fastapi.md
@@ -48,7 +48,7 @@ Next we'll create a FastAPI app and define a `/transactions/` POST endpoint:
 
 ```{literalinclude} ../../tests/fastapi/app.py
 :language: python
-:lines: 3,6,15-16,23-28
+:lines: 2-6,14-21,28-34
 ```
 
 ## Reading File Uploads
@@ -77,7 +77,7 @@ and the modified data in json format.
 
 ```{literalinclude} ../../tests/fastapi/app.py
 :language: python
-:lines: 7,30-38
+:lines: 37-44
 ```
 
 Pandera's {py:class}`~pandera.typing.fastapi.UploadFile` type is a subclass of FastAPI's
diff --git a/mypy.ini b/mypy.ini
index ea3653b..6c2d750 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -1,4 +1,5 @@
 [mypy]
+disable_error_code =annotation-unchecked
 ignore_missing_imports = True
 follow_imports = normal
 allow_redefinition = True
diff --git a/pandera/__init__.py b/pandera/__init__.py
index 8780f52..103edc3 100644
--- a/pandera/__init__.py
+++ b/pandera/__init__.py
@@ -75,7 +75,11 @@ from pandera.engines.pandas_engine import (
     pandas_version,
 )
 from pandera.schema_inference.pandas import infer_schema
-from pandera.version import __version__
+
+try:
+    from pandera._version import __version__
+except ImportError:
+    from pandera.version import __version__
 
 if platform.system() != "Windows":
     # pylint: disable=ungrouped-imports
diff --git a/pandera/api/base/model_components.py b/pandera/api/base/model_components.py
index 9c85b73..47ce8e4 100644
--- a/pandera/api/base/model_components.py
+++ b/pandera/api/base/model_components.py
@@ -136,8 +136,8 @@ class BaseCheckInfo:  # pylint:disable=too-few-public-methods
                 self.check_fn, "__name__", self.check_fn.__class__.__name__
             )
 
-        def _adapter(arg: Any) -> Union[bool, Iterable[bool]]:
-            return self.check_fn(model_cls, arg)
+        def _adapter(arg: Any, **kwargs) -> Union[bool, Iterable[bool]]:
+            return self.check_fn(model_cls, arg, **kwargs)
 
         return Check(_adapter, name=name, **self.check_kwargs)
 
diff --git a/pandera/api/pandas/model.py b/pandera/api/pandas/model.py
index 1887251..91658b5 100644
--- a/pandera/api/pandas/model.py
+++ b/pandera/api/pandas/model.py
@@ -207,14 +207,20 @@ class DataFrameModel(_DataFrameModel[pd.DataFrame, DataFrameSchema]):
         """
         schema = cls.to_schema()
         empty = pd.DataFrame(columns=schema.columns.keys()).astype(
-            {k: v.type for k, v in schema.dtypes.items()}
+            {k: v.type if v else None for k, v in schema.dtypes.items()}
         )
         table_schema = pd.io.json.build_table_schema(empty)
 
         def _field_json_schema(field):
             return {
                 "type": "array",
-                "items": {"type": field["type"]},
+                "items": {
+                    "type": (
+                        field["type"]
+                        if field["type"] != "any" or "extDtype" not in field
+                        else field["extDtype"]
+                    )
+                },
             }
 
         return {
diff --git a/pandera/api/polars/container.py b/pandera/api/polars/container.py
index 711b1b9..dafbd5f 100644
--- a/pandera/api/polars/container.py
+++ b/pandera/api/polars/container.py
@@ -3,8 +3,6 @@
 import warnings
 from typing import Optional, Type
 
-import polars as pl
-
 from pandera.api.dataframe.container import DataFrameSchema as _DataFrameSchema
 from pandera.api.polars.types import PolarsCheckObjects
 from pandera.api.polars.utils import get_validation_depth
@@ -55,12 +53,8 @@ class DataFrameSchema(_DataFrameSchema[PolarsCheckObjects]):
         if not get_config_context().validation_enabled:
             return check_obj
 
-        is_dataframe = isinstance(check_obj, pl.DataFrame)
         with config_context(validation_depth=get_validation_depth(check_obj)):
-            if is_dataframe:
-                # if validating a polars DataFrame, use the global config setting
-                check_obj = check_obj.lazy()
-
+            # if validating a polars DataFrame, use the global config setting
             output = self.get_backend(check_obj).validate(
                 check_obj=check_obj,
                 schema=self,
@@ -72,9 +66,6 @@ class DataFrameSchema(_DataFrameSchema[PolarsCheckObjects]):
                 inplace=inplace,
             )
 
-        if is_dataframe:
-            output = output.collect()
-
         return output
 
     @property
diff --git a/pandera/api/polars/types.py b/pandera/api/polars/types.py
index f038bcf..88c0452 100644
--- a/pandera/api/polars/types.py
+++ b/pandera/api/polars/types.py
@@ -1,6 +1,6 @@
 """Polars types."""
 
-from typing import NamedTuple, Optional, Union
+from typing import NamedTuple, Optional, Union, TypeVar
 
 import polars as pl
 
@@ -20,7 +20,7 @@ class CheckResult(NamedTuple):
 
 
 PolarsCheckObjects = Union[pl.LazyFrame, pl.DataFrame]
-
+PolarsFrame = TypeVar("PolarsFrame", pl.LazyFrame, pl.DataFrame)
 
 PolarsDtypeInputTypes = Union[
     str,
diff --git a/pandera/backends/pandas/container.py b/pandera/backends/pandas/container.py
index 081b7ca..b46142c 100644
--- a/pandera/backends/pandas/container.py
+++ b/pandera/backends/pandas/container.py
@@ -200,12 +200,12 @@ class DataFrameSchemaBackend(PandasSchemaBackend):
         check_passed = []
         # schema-component-level checks
         for schema_component in schema_components:
-            try:
-                # make sure the schema component mutations are reverted after
-                # validation
-                _orig_dtype = schema_component.dtype
-                _orig_coerce = schema_component.coerce
+            # make sure the schema component mutations are reverted after
+            # validation
+            _orig_dtype = schema_component.dtype
+            _orig_coerce = schema_component.coerce
 
+            try:
                 if schema.dtype is not None:
                     # override column dtype with dataframe dtype
                     schema_component.dtype = schema.dtype  # type: ignore
@@ -218,10 +218,6 @@ class DataFrameSchemaBackend(PandasSchemaBackend):
                     check_obj, lazy=lazy, inplace=True
                 )
 
-                # revert the schema component mutations
-                schema_component.dtype = _orig_dtype
-                schema_component.coerce = _orig_coerce
-
                 check_passed.append(is_table(result))
             except SchemaError as err:
                 check_results.append(
@@ -244,6 +240,11 @@ class DataFrameSchemaBackend(PandasSchemaBackend):
                         for schema_error in err.schema_errors
                     ]
                 )
+            finally:
+                # revert the schema component mutations
+                schema_component.dtype = _orig_dtype
+                schema_component.coerce = _orig_coerce
+
         assert all(check_passed)
         return check_results
 
diff --git a/pandera/backends/polars/base.py b/pandera/backends/polars/base.py
index 766efc9..9b756bc 100644
--- a/pandera/backends/polars/base.py
+++ b/pandera/backends/polars/base.py
@@ -7,7 +7,7 @@ from typing import Dict, List, Optional
 import polars as pl
 
 from pandera.api.base.error_handler import ErrorHandler
-from pandera.api.polars.types import CheckResult
+from pandera.api.polars.types import CheckResult, PolarsFrame
 from pandera.api.polars.utils import get_lazyframe_column_dtypes
 from pandera.backends.base import BaseSchemaBackend, CoreCheckResult
 from pandera.constants import CHECK_OUTPUT_KEY
@@ -34,24 +34,24 @@ class PolarsSchemaBackend(BaseSchemaBackend):
 
     def subsample(
         self,
-        check_obj: pl.LazyFrame,
+        check_obj: PolarsFrame,
         head: Optional[int] = None,
         tail: Optional[int] = None,
         sample: Optional[int] = None,
         random_state: Optional[int] = None,
-    ):
+    ) -> PolarsFrame:
         obj_subsample = []
         if head is not None:
             obj_subsample.append(check_obj.head(head))
         if tail is not None:
             obj_subsample.append(check_obj.tail(tail))
         if sample is not None:
-            obj_subsample.append(
-                # mypy is detecting a bug https://github.com/unionai-oss/pandera/issues/1912
-                check_obj.sample(  # type:ignore [attr-defined]
-                    sample, random_state=random_state
+            if isinstance(check_obj, pl.LazyFrame):
+                raise NotImplementedError(
+                    "Sample is not pl.LazyFrame inputs as polars does not implement LazyFrame.sample. "
+                    "Instead either pass a pl.DataFrame, or validate a subset using head/tail."
                 )
-            )
+            obj_subsample.append(check_obj.sample(sample, seed=random_state))
         return (
             check_obj
             if not obj_subsample
@@ -244,9 +244,9 @@ class PolarsSchemaBackend(BaseSchemaBackend):
 
     def drop_invalid_rows(
         self,
-        check_obj: pl.LazyFrame,
+        check_obj: PolarsFrame,
         error_handler: ErrorHandler,
-    ) -> pl.LazyFrame:
+    ) -> PolarsFrame:
         """Remove invalid elements in a check obj according to failures in caught by the error handler."""
         errors = error_handler.schema_errors
         check_outputs = pl.DataFrame(
diff --git a/pandera/backends/polars/container.py b/pandera/backends/polars/container.py
index 915a84a..c82d710 100644
--- a/pandera/backends/polars/container.py
+++ b/pandera/backends/polars/container.py
@@ -9,7 +9,7 @@ import polars as pl
 
 from pandera.api.base.error_handler import ErrorHandler
 from pandera.api.polars.container import DataFrameSchema
-from pandera.api.polars.types import PolarsData
+from pandera.api.polars.types import PolarsData, PolarsFrame
 from pandera.api.polars.utils import get_lazyframe_column_names
 from pandera.backends.base import ColumnInfo, CoreCheckResult
 from pandera.backends.polars.base import PolarsSchemaBackend
@@ -25,11 +25,25 @@ from pandera.utils import is_regex
 from pandera.validation_depth import validate_scope, validation_type
 
 
+def _to_lazy(df: PolarsFrame) -> pl.LazyFrame:
+    if isinstance(df, pl.DataFrame):
+        return df.lazy()
+    else:
+        return df
+
+
+def _to_frame_kind(lf: pl.LazyFrame, kind: type[PolarsFrame]) -> PolarsFrame:
+    if issubclass(kind, pl.DataFrame):
+        return lf.collect()
+    else:
+        return lf
+
+
 class DataFrameSchemaBackend(PolarsSchemaBackend):
     # pylint: disable=too-many-branches
     def validate(
         self,
-        check_obj: pl.LazyFrame,
+        check_obj: PolarsFrame,
         schema: DataFrameSchema,
         *,
         head: Optional[int] = None,
@@ -38,13 +52,16 @@ class DataFrameSchemaBackend(PolarsSchemaBackend):
         random_state: Optional[int] = None,
         lazy: bool = False,
         inplace: bool = False,
-    ):
+    ) -> PolarsFrame:
+        return_type = type(check_obj)
+        check_lf = _to_lazy(check_obj)  # parsers only accept lazyframe
+
         if inplace:
             warnings.warn("setting inplace=True will have no effect.")
 
         error_handler = ErrorHandler(lazy)
 
-        column_info = self.collect_column_info(check_obj, schema)
+        column_info = self.collect_column_info(check_lf, schema)
 
         if getattr(schema, "drop_invalid_rows", False) and not lazy:
             raise SchemaDefinitionError(
@@ -60,7 +77,7 @@ class DataFrameSchemaBackend(PolarsSchemaBackend):
 
         for parser, args in core_parsers:
             try:
-                check_obj = parser(check_obj, *args)
+                check_lf = parser(check_lf, *args)
             except SchemaError as exc:
                 error_handler.collect_error(
                     validation_type(exc.reason_code),
@@ -71,22 +88,29 @@ class DataFrameSchemaBackend(PolarsSchemaBackend):
                 error_handler.collect_errors(exc.schema_errors)
 
         components = self.collect_schema_components(
-            check_obj,
-            schema,
-            column_info,
+            check_lf, schema, column_info
         )
+        check_obj_parsed = _to_frame_kind(check_lf, return_type)
 
         # subsample the check object if head, tail, or sample are specified
-        sample = self.subsample(check_obj, head, tail, sample, random_state)
+        sample = self.subsample(
+            check_obj_parsed,
+            head,
+            tail,
+            sample,
+            random_state,
+        )
+        # all checks after subsampling are run on lazyframe
+        sample_lf = _to_lazy(sample)
 
         core_checks = [
-            (self.check_column_presence, (check_obj, schema, column_info)),
-            (self.check_column_values_are_unique, (sample, schema)),
+            (self.check_column_presence, (check_lf, schema, column_info)),
+            (self.check_column_values_are_unique, (sample_lf, schema)),
             (
                 self.run_schema_component_checks,
-                (sample, schema, components, lazy),
+                (sample_lf, schema, components, lazy),
             ),
-            (self.run_checks, (sample, schema)),
+            (self.run_checks, (sample_lf, schema)),
         ]
 
         for check, args in core_checks:
@@ -104,7 +128,7 @@ class DataFrameSchemaBackend(PolarsSchemaBackend):
                 else:
                     error = SchemaError(
                         schema,
-                        data=check_obj,
+                        data=check_lf,
                         message=result.message,
                         failure_cases=result.failure_cases,
                         check=result.check,
@@ -121,15 +145,17 @@ class DataFrameSchemaBackend(PolarsSchemaBackend):
 
         if error_handler.collected_errors:
             if getattr(schema, "drop_invalid_rows", False):
-                check_obj = self.drop_invalid_rows(check_obj, error_handler)
+                check_obj_parsed = self.drop_invalid_rows(
+                    check_obj_parsed, error_handler
+                )
             else:
                 raise SchemaErrors(
                     schema=schema,
                     schema_errors=error_handler.schema_errors,
-                    data=check_obj,
+                    data=check_obj_parsed,
                 )
 
-        return check_obj
+        return check_obj_parsed
 
     @validate_scope(scope=ValidationScope.DATA)
     def run_checks(
diff --git a/pandera/backends/polars/register.py b/pandera/backends/polars/register.py
index 34473e5..e24a0b5 100644
--- a/pandera/backends/polars/register.py
+++ b/pandera/backends/polars/register.py
@@ -26,5 +26,6 @@ def register_polars_backends(
     from pandera.backends.polars.container import DataFrameSchemaBackend
 
     DataFrameSchema.register_backend(pl.LazyFrame, DataFrameSchemaBackend)
+    DataFrameSchema.register_backend(pl.DataFrame, DataFrameSchemaBackend)
     Column.register_backend(pl.LazyFrame, ColumnBackend)
     Check.register_backend(pl.LazyFrame, PolarsCheckBackend)
diff --git a/pandera/engines/polars_engine.py b/pandera/engines/polars_engine.py
index 9ba452c..6d4b8c5 100644
--- a/pandera/engines/polars_engine.py
+++ b/pandera/engines/polars_engine.py
@@ -583,9 +583,19 @@ class Array(DataType):
 
     @classmethod
     def from_parametrized_dtype(cls, polars_dtype: pl.Array):
+        shape = polars_dtype.shape
+
+        if (
+            isinstance(shape, tuple)
+            and isinstance(shape[0], int)
+            and len(shape) > 1
+        ):
+            offset = len(shape) - 1
+            shape = shape[offset:]
+
         return cls(
             inner=polars_dtype.inner,
-            shape=polars_dtype.shape,
+            shape=shape,
         )
 
 
diff --git a/pandera/typing/pandas.py b/pandera/typing/pandas.py
index 91e7e92..1786a65 100644
--- a/pandera/typing/pandas.py
+++ b/pandera/typing/pandas.py
@@ -198,24 +198,34 @@ class DataFrame(DataFrameBase, pd.DataFrame, Generic[T]):
                 "boolean": core_schema.bool_schema(),
                 "datetime": core_schema.datetime_schema(),
             }
-            return core_schema.no_info_plain_validator_function(
-                functools.partial(
-                    cls.pydantic_validate,
-                    schema_model=schema_model,
-                ),
-                json_schema_input_schema=core_schema.list_schema(
-                    core_schema.typed_dict_schema(
-                        {
-                            key: core_schema.typed_dict_field(
-                                type_map[
-                                    schema_json_columns[key]["items"]["type"]
-                                ]
-                            )
-                            for key in schema.columns.keys()
-                        },
-                    )
-                ),
+            function = functools.partial(
+                cls.pydantic_validate,
+                schema_model=schema_model,
             )
+            json_schema_input_schema = core_schema.list_schema(
+                core_schema.typed_dict_schema(
+                    {
+                        key: core_schema.typed_dict_field(
+                            type_map[schema_json_columns[key]["items"]["type"]]
+                        )
+                        for key in schema.columns.keys()
+                    },
+                )
+            )
+            try:
+                # json schema input schema is only available in
+                # pydantic_core >=2.30.0
+                return core_schema.no_info_plain_validator_function(
+                    function,
+                    json_schema_input_schema=json_schema_input_schema,
+                    serialization=core_schema.plain_serializer_function_ser_schema(
+                        function=lambda df: df,
+                        info_arg=False,
+                        return_schema=json_schema_input_schema,
+                    ),
+                )
+            except TypeError:
+                return core_schema.no_info_plain_validator_function(function)
 
     else:
 


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 7fcef2c..4b4854f 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -5,7 +5,8 @@ import os
 import re
 import runpy
 from copy import deepcopy
-from typing import Any, Generic, Iterable, List, Optional, TypeVar
+from enum import Enum
+from typing import Any, Generic, Iterable, List, Optional, TypeVar, Type
 
 import numpy as np
 import pandas as pd
@@ -640,6 +641,50 @@ def test_dataframe_check() -> None:
     assert len(e.value.message["DATA"]) == 1
 
 
+def test_dataframe_check_passthrough_kwargs() -> None:
+    """Test dataframe checks with passthrough kwargs."""
+
+    class MockTypesOne(Enum):
+        ONE = 1
+        TWO = 2
+        THREE = 3
+
+    class MockTypesTwo(Enum):
+        AVG = "avg"
+        SUM = "sum"
+        MEAN = "mean"
+
+    class Base(pa.DataFrameModel):
+        a: Series[int]
+        b: Series[str]
+
+        @pa.dataframe_check(column="a", enum=MockTypesOne)
+        @classmethod
+        def type_id_valid(
+            cls, df: pd.DataFrame, column: str, enum: Type[Enum]
+        ) -> Iterable[bool]:
+            return cls._field_in_enum(df, column, enum)  # type: ignore
+
+        @pa.dataframe_check(column="b", enum=MockTypesTwo)
+        @classmethod
+        def calc_type_valid(
+            cls, df: pd.DataFrame, column: str, enum: Type[Enum]
+        ) -> Iterable[bool]:
+            return cls._field_in_enum(df, column, enum)  # type: ignore
+
+        @classmethod
+        def _field_in_enum(
+            cls, df: pd.DataFrame, column: str, enum: Type[Enum]
+        ) -> Iterable[bool]:
+            return df[column].isin([member.value for member in enum])  # type: ignore
+
+    schema = Base.to_schema()
+    assert len(schema.checks) == 2
+
+    df = pd.DataFrame({"a": [1, 2], "b": ["avg", "mean"]})
+    schema.validate(df, lazy=True)
+
+
 def test_registered_dataframe_checks(
     extra_registered_checks: None,  # pylint: disable=unused-argument
 ) -> None:
@@ -1556,3 +1601,23 @@ def test_empty() -> None:
     df = Schema.empty()
     assert df.empty
     assert Schema.validate(df).empty  # type: ignore [attr-defined]
+
+
+def test_model_with_pydantic_base_model_with_df_init():
+    """
+    Test that a dataframe can be initialized within a pydantic base model that
+    defines the dataframe model as a class attribute.
+    """
+    from pydantic import BaseModel
+
+    # pylint: disable=unused-variable
+    class PydanticModel(BaseModel):
+
+        class PanderaDataFrameModel(pa.DataFrameModel):
+            """The DF we use to represent code/label/description associations."""
+
+            field: pa.typing.Series[Any]
+
+        df: pa.typing.DataFrame[PanderaDataFrameModel] = pd.DataFrame(
+            {"field": [1, 2, 3]}
+        )
diff --git a/tests/core/test_pydantic.py b/tests/core/test_pydantic.py
index cf1755b..dc7c3c1 100644
--- a/tests/core/test_pydantic.py
+++ b/tests/core/test_pydantic.py
@@ -180,3 +180,31 @@ def test_invalid_seriesschema():
 
     with pytest.raises(TypeError if PYDANTIC_V2 else ValidationError):
         SeriesSchemaPydantic(pa_index="1")
+
+
+@pytest.mark.parametrize(
+    "col_type,dtype,item",
+    [
+        (pa.STRING, "string", "hello"),
+        (pa.UINT8, "UInt8", 1),
+        (pa.INT8, "Int8", 1),
+        (pa.BOOL, "boolean", True),
+    ],
+)
+def test_model_with_extensiondtype_column(col_type, dtype, item):
+    """Test that a model with an external dtype is recognized by pydantic."""
+
+    class ExtensionDtypeModel(pa.DataFrameModel):
+        a: Series[col_type]
+
+    class PydanticModel(BaseModel):
+        df: DataFrame[ExtensionDtypeModel]
+
+    assert isinstance(
+        PydanticModel(
+            df=DataFrame[ExtensionDtypeModel](
+                pd.DataFrame({"a": [item]}, dtype=dtype)
+            )
+        ),
+        PydanticModel,
+    )
diff --git a/tests/core/test_schemas.py b/tests/core/test_schemas.py
index b3ac4eb..cca335c 100644
--- a/tests/core/test_schemas.py
+++ b/tests/core/test_schemas.py
@@ -2477,6 +2477,33 @@ def test_schema_coerce_with_regex() -> None:
     assert isinstance(schema_with_regex.validate(df), pd.DataFrame)
 
 
+def test_schema_coerce_preserve_value():
+    """Test that coercing an invalid data preserves the original coerce value."""
+
+    schema = DataFrameSchema(
+        {
+            "chr": Column(
+                str, checks=Check.str_length(min_value=1), coerce=True
+            ),
+            "start": Column(int, checks=Check.ge(0)),
+        }
+    )
+    assert schema.columns["chr"].coerce
+
+    schema.validate(pd.DataFrame({"chr": ["chr1", "chr2"], "start": [0, 10]}))
+    assert schema.columns["chr"].coerce
+
+    try:
+        schema.validate(pd.DataFrame({"chr": ["", "chr1"], "start": [0, 10]}))
+        raise AssertionError(
+            "Dataframe should fail validation as str_length constraint not met"
+        )
+    except errors.SchemaError:
+        ...
+
+    assert schema.columns["chr"].coerce
+
+
 @pytest.mark.parametrize(
     "schema, obj, expected_obj",
     [
diff --git a/tests/polars/test_polars_container.py b/tests/polars/test_polars_container.py
index edcaf10..3d79051 100644
--- a/tests/polars/test_polars_container.py
+++ b/tests/polars/test_polars_container.py
@@ -487,6 +487,17 @@ def test_ordered(ldf_basic, ldf_schema_basic):
         invalid_order.pipe(ldf_schema_basic.validate).collect()
 
 
+def test_sample_dataframe_schema(df_basic, ldf_basic, ldf_schema_basic):
+
+    with pytest.raises(NotImplementedError):
+        ldf_schema_basic.validate(ldf_basic, sample=1, random_state=1)
+
+    validated_data = ldf_schema_basic.validate(
+        df_basic, sample=1, random_state=1
+    )
+    assert validated_data.equals(df_basic)
+
+
 @pytest.mark.parametrize("arg", ["exclude_first", "exclude_last"])
 def test_report_duplicates(arg):
     with pytest.warns(
diff --git a/tests/polars/test_polars_dtypes.py b/tests/polars/test_polars_dtypes.py
index 34cafe0..28a84f9 100644
--- a/tests/polars/test_polars_dtypes.py
+++ b/tests/polars/test_polars_dtypes.py
@@ -450,6 +450,37 @@ def test_polars_nested_dtypes_try_coercion(
         assert exc.failure_cases.select(col).equals(data.collect())
 
 
+@pytest.mark.parametrize(
+    "array",
+    [
+        pl.Array(pl.Int64(), (2, 2)),
+        pl.Array(pl.Int64(), (2, 2, 2)),
+        pl.Array(pl.Int64(), (2, 2, 2, 2)),
+    ],
+)
+def test_polars_nested_dtypes_shape(array):
+    pandera_dtype = pe.Engine.dtype(array)
+
+    assert len(array.shape) == len(pandera_dtype.type.shape)
+    assert array.shape == pandera_dtype.type.shape
+
+
+@pytest.mark.parametrize(
+    "dtype, shape",
+    [
+        (pl.Int64(), (2, 2)),
+        (pl.Int64(), (2, 2, 2)),
+        (pl.Int64(), (2, 2, 2, 2)),
+    ],
+)
+def test_polars_from_parametrized_nested_dtype(dtype, shape):
+    polars_array_type = pl.Array(dtype, shape=shape)
+    pandera_dtype = pe.Array.from_parametrized_dtype(polars_array_type)
+
+    assert pandera_dtype.type.shape == polars_array_type.shape
+    assert pandera_dtype.type.shape == shape
+
+
 @pytest.mark.parametrize(
     "dtype",
     [