feat: Support Dataframes as Fields in Pydantic Models (#148)

Co-authored-by: Andreas Albert <[email protected]>

Author: ytausch

dataframely/_base_schema.py

@@ -5,10 +5,10 @@
 
 import sys
 import textwrap
-from abc import ABCMeta
+from abc import ABCMeta, abstractmethod
 from copy import copy
 from dataclasses import dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import polars as pl
 
@@ -21,6 +21,10 @@
 else:
     from typing_extensions import Self
 
+
+if TYPE_CHECKING:
+    from ._typing import DataFrame
+
 _COLUMN_ATTR = "__dataframely_columns__"
 _RULE_ATTR = "__dataframely_rules__"
 
@@ -198,11 +202,46 @@ def columns(cls) -> dict[str, Column]:
             columns[name]._name = name
         return columns
 
+    @classmethod
+    @abstractmethod
+    def polars_schema(cls) -> pl.Schema:
+        """Obtain the polars schema for this schema.
+
+        Returns:
+            A :mod:`polars` schema that mirrors the schema defined by this class.
+        """
+
     @classmethod
     def primary_keys(cls) -> list[str]:
         """The primary key columns in this schema (possibly empty)."""
         return _primary_keys(cls.columns())
 
+    @classmethod
+    @abstractmethod
+    def validate(
+        cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False
+    ) -> DataFrame[Self]:
+        """Validate that a data frame satisfies the schema.
+
+        Args:
+            df: The data frame to validate.
+            cast: Whether columns with a wrong data type in the input data frame are
+                cast to the schema's defined data type if possible.
+
+        Returns:
+            The (collected) input data frame, wrapped in a generic version of the
+            input's data frame type to reflect schema adherence. The data frame is
+            guaranteed to maintain its order.
+
+        Raises:
+            ValidationError: If the input data frame does not satisfy the schema
+                definition.
+
+        Note:
+            This method _always_ collects the input data frame in order to raise
+            potential validation errors.
+        """
+
     @classmethod
     def _validation_rules(cls, *, with_cast: bool) -> dict[str, Rule]:
         return _build_rules(

dataframely/_compat.py

@@ -54,6 +54,19 @@ class Dialect:  # type: ignore # noqa: N801
 except ImportError:  # pragma: no cover
     pa = _DummyModule("pyarrow")
 
+
+# -------------------------------------- PYDANTIC ------------------------------------ #
+
+try:
+    import pydantic
+except ImportError:  # pragma: no cover
+    pydantic = _DummyModule("pydantic")  # type: ignore
+
+try:
+    from pydantic_core import core_schema as pydantic_core_schema  # pragma: no cover
+except ImportError:
+    pydantic_core_schema = _DummyModule("pydantic_core_schema")  # type: ignore
+
 # ------------------------------------------------------------------------------------ #
 
 __all__ = [
@@ -64,4 +77,6 @@ class Dialect:  # type: ignore # noqa: N801
     "pa",
     "MSDialect_pyodbc",
     "PGDialect_psycopg2",
+    "pydantic",
+    "pydantic_core_schema",
 ]

dataframely/_pydantic.py

@@ -0,0 +1,115 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING, Literal, TypeVar, get_args, get_origin, overload
+
+import polars as pl
+
+from ._base_schema import BaseSchema
+from ._compat import pydantic, pydantic_core_schema
+from .exc import ValidationError
+
+if TYPE_CHECKING:
+    from ._typing import DataFrame, LazyFrame
+
+
+_S = TypeVar("_S", bound=BaseSchema)
+
+
+def _dict_to_df(schema_type: type[BaseSchema], data: dict) -> pl.DataFrame:
+    return pl.from_dict(
+        data,
+        schema=schema_type.polars_schema(),
+    )
+
+
+def _validate_df_schema(schema_type: type[_S], df: pl.DataFrame) -> DataFrame[_S]:
+    try:
+        return schema_type.validate(df, cast=False)
+    except ValidationError as e:
+        raise ValueError("DataFrame violates schema") from e
+
+
+def _serialize_df(df: pl.DataFrame) -> dict:
+    return df.to_dict(as_series=False)
+
+
+@overload
+def get_pydantic_core_schema(
+    source_type: type[DataFrame],
+    _handler: pydantic.GetCoreSchemaHandler,
+    lazy: Literal[False],
+) -> pydantic_core_schema.CoreSchema: ...
+
+
+@overload
+def get_pydantic_core_schema(
+    source_type: type[LazyFrame],
+    _handler: pydantic.GetCoreSchemaHandler,
+    lazy: Literal[True],
+) -> pydantic_core_schema.CoreSchema: ...
+
+
+def get_pydantic_core_schema(
+    source_type: type[DataFrame | LazyFrame],
+    _handler: pydantic.GetCoreSchemaHandler,
+    lazy: bool,
+) -> pydantic_core_schema.CoreSchema:
+    # https://docs.pydantic.dev/2.11/concepts/types/#handling-custom-generic-classes
+    origin = get_origin(source_type)
+    if origin is None:
+        # used as `x: dy.DataFrame` without schema
+        raise TypeError("DataFrame must be parametrized with a schema")
+
+    schema_type: type[BaseSchema] = get_args(source_type)[0]
+
+    # accept a DataFrame, a LazyFrame, or a dict that is converted to a DataFrame
+    # (-> output: DataFrame or LazyFrame)
+    polars_schema = pydantic_core_schema.union_schema(
+        [
+            pydantic_core_schema.is_instance_schema(pl.DataFrame),
+            pydantic_core_schema.is_instance_schema(pl.LazyFrame),
+            pydantic_core_schema.chain_schema(
+                [
+                    pydantic_core_schema.dict_schema(),
+                    pydantic_core_schema.no_info_plain_validator_function(
+                        partial(_dict_to_df, schema_type)
+                    ),
+                ]
+            ),
+        ]
+    )
+
+    to_lazy_schema = []
+    if lazy:
+        # If the Pydantic field type is LazyFrame, add a step to convert
+        # the model back to a LazyFrame.
+        to_lazy_schema.append(
+            pydantic_core_schema.no_info_plain_validator_function(
+                lambda df: df.lazy(),
+            )
+        )
+
+    return pydantic_core_schema.chain_schema(
+        [
+            polars_schema,
+            pydantic_core_schema.no_info_plain_validator_function(
+                partial(_validate_df_schema, schema_type)
+            ),
+            *to_lazy_schema,
+        ],
+        serialization=pydantic_core_schema.plain_serializer_function_ser_schema(
+            _serialize_df
+        ),
+    )
+
+
+def get_pydantic_json_schema(
+    handler: pydantic.GetJsonSchemaHandler,
+) -> pydantic.json_schema.JsonSchemaValue:
+    from pydantic_core import core_schema
+
+    # This could be made more sophisticated by actually reflecting the schema.
+    return handler(core_schema.dict_schema())

dataframely/_typing.py

@@ -9,6 +9,8 @@
 import polars as pl
 
 from ._base_schema import BaseSchema
+from ._compat import pydantic, pydantic_core_schema
+from ._pydantic import get_pydantic_core_schema, get_pydantic_json_schema
 
 S = TypeVar("S", bound=BaseSchema, covariant=True)
 
@@ -70,6 +72,20 @@ def set_sorted(self, *args: Any, **kwargs: Any) -> DataFrame[S]:
     def shrink_to_fit(self, *args: Any, **kwargs: Any) -> DataFrame[S]:
         raise NotImplementedError  # pragma: no cover
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> pydantic_core_schema.CoreSchema:
+        return get_pydantic_core_schema(source_type, handler, lazy=False)
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        _core_schema: pydantic_core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> pydantic.json_schema.JsonSchemaValue:
+        return get_pydantic_json_schema(handler)
+
 
 class LazyFrame(pl.LazyFrame, Generic[S]):
     """Generic wrapper around a :class:`polars.LazyFrame` to attach schema information.
@@ -113,3 +129,17 @@ def pipe(
     @inherit_signature(pl.LazyFrame.set_sorted)
     def set_sorted(self, *args: Any, **kwargs: Any) -> LazyFrame[S]:
         raise NotImplementedError  # pragma: no cover
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> pydantic_core_schema.CoreSchema:
+        return get_pydantic_core_schema(source_type, handler, lazy=True)
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        _core_schema: pydantic_core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> pydantic.json_schema.JsonSchemaValue:
+        return get_pydantic_json_schema(handler)

dataframely/schema.py

@@ -414,26 +414,6 @@ def _sampling_overrides(cls) -> dict[str, pl.Expr]:
     def validate(
         cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False
     ) -> DataFrame[Self]:
-        """Validate that a data frame satisfies the schema.
-
-        Args:
-            df: The data frame to validate.
-            cast: Whether columns with a wrong data type in the input data frame are
-                cast to the schema's defined data type if possible.
-
-        Returns:
-            The (collected) input data frame, wrapped in a generic version of the
-            input's data frame type to reflect schema adherence. The data frame is
-            guaranteed to maintain its order.
-
-        Raises:
-            ValidationError: If the input data frame does not satisfy the schema
-                definition.
-
-        Note:
-            This method _always_ collects the input data frame in order to raise
-            potential validation errors.
-        """
         # We can dispatch to the `filter` method and raise an error if any row cannot
         # be validated
         df_valid, failures = cls.filter(df, cast=cast)
@@ -1118,11 +1098,6 @@ def _validate_if_needed(
 
     @classmethod
     def polars_schema(cls) -> pl.Schema:
-        """Obtain the polars schema for this schema.
-
-        Returns:
-            A :mod:`polars` schema that mirrors the schema defined by this class.
-        """
         return pl.Schema({name: col.dtype for name, col in cls.columns().items()})
 
     @classmethod

docs/sites/quickstart.rst

@@ -253,6 +253,7 @@ Lastly, ``dataframely`` schemas can be used to integrate with external tools:
 - ``HouseSchema.create_empty()`` creates an empty ``dy.DataFrame[HouseSchema]`` that can be used for testing
 - ``HouseSchema.sql_schema()`` provides a list of `sqlalchemy <https://www.sqlalchemy.org>`_ columns that can be used to create SQL tables using types and constraints in line with the schema
 - ``HouseSchema.pyarrow_schema()`` provides a `pyarrow <https://arrow.apache.org/docs/python/index.html>`_ schema with appropriate column dtypes and nullability information
+- You can use ``dy.DataFrame[HouseSchema]`` (or the ``LazyFrame`` equivalent) as fields in `pydantic <https://pydantic.dev>`_ models, including support for validation and serialization. Integration with pydantic is unstable.
 
 
 Outlook