diff --git a/.github/workflows/ci-tests.yml b/.github/workflows/ci-tests.yml index 0e95b7de8..5bb0a97c5 100644 --- a/.github/workflows/ci-tests.yml +++ b/.github/workflows/ci-tests.yml @@ -335,6 +335,7 @@ jobs: - modin-ray - ibis - xarray + - pyarrow include: - extra: polars polars-version: "0.20.0" diff --git a/docs/source/index.md b/docs/source/index.md index 3b68642dc..345d41290 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -159,6 +159,7 @@ pip install 'pandera[modin-dask]' # validate modin dataframes with dask pip install 'pandera[geopandas]' # validate geopandas geodataframes pip install 'pandera[polars]' # validate polars dataframes pip install 'pandera[ibis]' # validate ibis tables +pip install 'pandera[pyarrow]' # validate pyarrow tables pip install 'pandera[xarray]' # validate xarray data structures pip install 'pandera[narwhals]' # use the Narwhals-powered backend ``` diff --git a/docs/source/pyarrow.md b/docs/source/pyarrow.md new file mode 100644 index 000000000..3154be239 --- /dev/null +++ b/docs/source/pyarrow.md @@ -0,0 +1,138 @@ +--- +file_format: mystnb +--- + +% pandera documentation for pyarrow + +```{currentmodule} pandera.pyarrow +``` + +(pyarrow)= + +# Data Validation with PyArrow + +*new in 0.33.0* + +[PyArrow](https://arrow.apache.org/docs/python/) tables are a common exchange +format between systems, and `pyarrow.Table` carries its own schema. Pandera +supports validating them directly so you get dtype checking, value checks and +class-based models on top of that schema. + +## Installation + +```bash +pip install 'pandera[pyarrow]' +``` + +Validation is performed by pandera's +{ref}`narwhals backend `, so `narwhals` is installed +alongside `pyarrow`. + +## `DataFrameSchema` + +Import the pyarrow entry point and define a schema as you would for any other +backend: + +```{code-cell} python +import pyarrow +import pandera.pyarrow as pa + +schema = pa.DataFrameSchema( + { + "state": pa.Column(str), + "city": pa.Column(str), + "price": pa.Column(int, pa.Check.in_range(5, 20)), + } +) + +table = pyarrow.table( + { + "state": ["FL", "FL", "CA", "CA"], + "city": ["Orlando", "Miami", "Los Angeles", "San Francisco"], + "price": [8, 12, 10, 16], + } +) + +schema.validate(table) +``` + +`validate` returns a `pyarrow.Table`, so schemas drop into an existing pipeline +without changing types. + +## `DataFrameModel` + +```{code-cell} python +class Schema(pa.DataFrameModel): + state: str + city: str + price: int = pa.Field(in_range={"min_value": 5, "max_value": 20}) + + +Schema.validate(table) +``` + +Annotate function signatures with `pandera.typing.pyarrow.Table` and use +{func}`~pandera.decorators.check_types` to validate inputs and outputs: + +```{code-cell} python +from pandera.typing.pyarrow import Table + + +@pa.check_types +def transform(df: Table[Schema]) -> Table[Schema]: + return df + + +transform(table) +``` + +## Supported data types + +Columns accept native pyarrow types, python builtins, and their string +aliases — all resolve to the same underlying pandera datatype: + +```{code-cell} python +pa.DataFrameSchema( + { + "a": pa.Column(pyarrow.int64()), + "b": pa.Column(int), + "c": pa.Column("int64"), + } +) +``` + +Parametrized pyarrow types such as `pyarrow.timestamp("us")`, +`pyarrow.decimal128(10, 2)` and `pyarrow.list_(pyarrow.int32())` are supported. + +## Custom checks + +Check functions receive a {class}`~pandera.api.pyarrow.types.PyArrowData` +container holding the native table and the column key, mirroring `PolarsData` +and `IbisData` on the other backends: + +```{code-cell} python +import pyarrow.compute as pc + +schema = pa.DataFrameSchema( + {"price": pa.Column(int, pa.Check(lambda data: pc.greater(data.table[data.key], 0)))} +) +schema.validate(table) +``` + +A check function taking two positional arguments receives the native table and +the key directly, i.e. `check_fn(table, key)`. + +## Validation depth + +Unlike a `polars.LazyFrame` or an `ibis.Table`, a `pyarrow.Table` is always +fully materialized in memory, so pandera runs both schema-level and data-level +checks by default. Set `PANDERA_VALIDATION_DEPTH=SCHEMA_ONLY` (or use +{func}`~pandera.config.config_context`) to skip data-level checks. + +## Limitations + +- **`coerce=True` is not applied.** Coercion is not yet implemented in the + narwhals backend that serves pyarrow; a `SchemaWarning` is emitted and a + dtype mismatch is reported as a `WRONG_DATATYPE` error instead. +- **Data synthesis strategies are not supported**, matching the polars and ibis + backends. diff --git a/docs/source/supported_libraries.md b/docs/source/supported_libraries.md index 2ffc9b72d..c0d566b5f 100644 --- a/docs/source/supported_libraries.md +++ b/docs/source/supported_libraries.md @@ -23,6 +23,8 @@ Pandera supports validation of the following DataFrame libraries: - Validate Polars dataframes. Polars is a blazingly fast dataframe library. * - {ref}`Ibis ` - Validate Ibis tables. Ibis is the portable Python dataframe library. +* - {ref}`PyArrow ` + - Validate PyArrow tables. Arrow is the in-memory columnar exchange format. * - {ref}`PySpark SQL ` - A data processing library for large-scale data. ::: @@ -33,6 +35,7 @@ Pandera supports validation of the following DataFrame libraries: Polars Ibis +PyArrow PySpark SQL ``` diff --git a/noxfile.py b/noxfile.py index 8a4ee5a0d..64ad77d0a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -221,6 +221,7 @@ def _testing_requirements( "ibis", "xarray", "narwhals", # TEST-03: narwhals backend runs with polars+ibis co-installed + "pyarrow", # pyarrow.Table validation, served by the narwhals backends } for extra in OPTIONAL_DEPENDENCIES: if extra == "pandas": diff --git a/pandera/api/pyarrow/__init__.py b/pandera/api/pyarrow/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pandera/api/pyarrow/components.py b/pandera/api/pyarrow/components.py new file mode 100644 index 000000000..c8803b5c4 --- /dev/null +++ b/pandera/api/pyarrow/components.py @@ -0,0 +1,174 @@ +"""Core PyArrow schema component specifications.""" + +from __future__ import annotations + +import logging +from typing import Any + +from pandera.api.base.types import CheckList +from pandera.api.dataframe.components import ComponentSchema +from pandera.api.pyarrow.types import ( + PyArrowCheckObjects, + PyArrowDtypeInputTypes, +) +from pandera.api.pyarrow.utils import resolve_dtype +from pandera.backends.pyarrow.register import register_pyarrow_backends +from pandera.utils import is_regex + +logger = logging.getLogger(__name__) + + +class Column(ComponentSchema[PyArrowCheckObjects]): + """Validate types and properties of pyarrow table columns.""" + + def __init__( + self, + dtype: PyArrowDtypeInputTypes | None = None, + checks: CheckList | None = None, + nullable: bool = False, + unique: bool = False, + coerce: bool = False, + required: bool = True, + name: str | None = None, + regex: bool = False, + title: str | None = None, + description: str | None = None, + default: Any | None = None, + metadata: dict | None = None, + drop_invalid_rows: bool = False, + **column_kwargs, + ) -> None: + """Create column validator object. + + :param dtype: datatype of the column. Accepts native ``pyarrow`` + types (e.g. ``pyarrow.int64()``), narwhals dtypes, the pandera + abstract datatypes, supported python builtins (``int``, ``float``, + ``str``, ``bool``) and their string aliases. + :param checks: checks to verify validity of the column + :param nullable: Whether or not column can contain null values. + :param unique: whether column values should be unique + :param coerce: If True, when schema.validate is called the column will + be coerced into the specified dtype. This has no effect on columns + where ``dtype=None``. + :param required: Whether or not column is allowed to be missing + :param name: column name in the table to validate. Names in the format + '^{regex_pattern}$' are treated as regular expressions. During + validation, this schema will be applied to any columns matching + this pattern. + :param regex: whether the ``name`` attribute should be treated as a + regex pattern to apply to multiple columns in a table. + :param title: A human-readable label for the column. + :param description: An arbitrary textual description of the column. + :param default: The default value for missing values in the column. + :param metadata: An optional key value data. + :param drop_invalid_rows: if True, drop invalid rows on validation. + + :raises SchemaInitError: if impossible to build schema from parameters + + :example: + + >>> import pyarrow + >>> import pandera.pyarrow as pa + >>> + >>> schema = pa.DataFrameSchema({"column": pa.Column(str)}) + >>> schema.validate(pyarrow.table({"column": ["foo", "bar"]})) + pyarrow.Table + column: string + ---- + column: [["foo","bar"]] + """ + super().__init__( + dtype=dtype, + checks=checks, + nullable=nullable, + unique=unique, + coerce=coerce, + name=name, + title=title, + description=description, + default=default, + metadata=metadata, + drop_invalid_rows=drop_invalid_rows, + **column_kwargs, + ) + self.required = required + self.regex = regex + self.name = name + + self.set_regex() + + @staticmethod + def register_default_backends(check_obj_cls: type): + register_pyarrow_backends() + + @property + def dtype(self): + return self._dtype + + @dtype.setter + def dtype(self, value) -> None: + self._dtype = resolve_dtype(value) + + @property + def selector(self): + if self.name is not None and not is_regex(self.name) and self.regex: + return f"^{self.name}$" + return self.name + + def set_regex(self): + if self.name is None: + return + + if is_regex(self.name) and not self.regex: + logger.info( + f"Column schema '{self.name}' is a regex expression. " + "Setting regex=True." + ) + self.regex = True + + def set_name(self, name: str): + """Set or modify the name of a column object. + + :param str name: the name of the column object + """ + self.name = name + self.set_regex() + return self + + @property + def properties(self) -> dict[str, Any]: + """Get column properties.""" + return { + "dtype": self.dtype, + "parsers": self.parsers, + "checks": self.checks, + "nullable": self.nullable, + "unique": self.unique, + "report_duplicates": self.report_duplicates, + "coerce": self.coerce, + "required": self.required, + "name": self.name, + "regex": self.regex, + "title": self.title, + "description": self.description, + "default": self.default, + "metadata": self.metadata, + } + + def strategy(self, *, size=None): + """Data synthesis is not supported for pyarrow schemas.""" + raise NotImplementedError( + "Data synthesis is not supported with pyarrow schemas." + ) + + def strategy_component(self): + """Data synthesis is not supported for pyarrow schemas.""" + raise NotImplementedError( + "Data synthesis is not supported with pyarrow schemas." + ) + + def example(self, size=None): + """Data synthesis is not supported for pyarrow schemas.""" + raise NotImplementedError( + "Data synthesis is not supported with pyarrow schemas." + ) diff --git a/pandera/api/pyarrow/container.py b/pandera/api/pyarrow/container.py new file mode 100644 index 000000000..6a697de1d --- /dev/null +++ b/pandera/api/pyarrow/container.py @@ -0,0 +1,113 @@ +"""Core PyArrow table container specification.""" + +from __future__ import annotations + +import warnings + +import pyarrow as pa + +from pandera.api.dataframe.container import DataFrameSchema as _DataFrameSchema +from pandera.api.pyarrow.utils import get_validation_depth, resolve_dtype +from pandera.backends.pyarrow.register import register_pyarrow_backends +from pandera.config import config_context, get_config_context + + +class DataFrameSchema(_DataFrameSchema[pa.Table]): + """A lightweight pyarrow Table validator.""" + + def _validate_attributes(self): + super()._validate_attributes() + + if self.unique_column_names: + warnings.warn( + "unique_column_names=True will have no effect on validation " + "since pyarrow Tables do not support duplicate column names." + ) + + if self.report_duplicates != "all": + warnings.warn( + "Setting report_duplicates to 'exclude_first' or " + "'exclude_last' will have no effect on validation. With the " + "pyarrow backend, all duplicate values will be reported." + ) + + @staticmethod + def register_default_backends(check_obj_cls: type): + register_pyarrow_backends() + + def validate( + self, + check_obj: pa.Table, + head: int | None = None, + tail: int | None = None, + sample: int | None = None, + random_state: int | None = None, + lazy: bool = False, + inplace: bool = False, + ) -> pa.Table: + """Validate a pyarrow Table against the schema. + + :param check_obj: the ``pyarrow.Table`` to be validated. + :param head: validate the first n rows. Rows overlapping with ``tail`` + or ``sample`` are de-duplicated. + :param tail: validate the last n rows. Rows overlapping with ``head`` + or ``sample`` are de-duplicated. + :param sample: validate a random sample of n rows. Rows overlapping + with ``head`` or ``tail`` are de-duplicated. + :param random_state: random seed for the ``sample`` argument. + :param lazy: if True, lazily evaluates the table against all validation + checks and raises a ``SchemaErrors``. Otherwise, raise + ``SchemaError`` as soon as one occurs. + :param inplace: has no effect; pyarrow Tables are immutable. + :returns: validated ``pyarrow.Table`` + + :raises SchemaError: when ``check_obj`` violates built-in or custom + checks. + + :example: + + >>> import pyarrow + >>> import pandera.pyarrow as pa + >>> + >>> schema = pa.DataFrameSchema({ + ... "probability": pa.Column(float, pa.Check.le(1.0)), + ... }) + >>> schema.validate(pyarrow.table({"probability": [0.1, 0.4]})) + pyarrow.Table + probability: double + ---- + probability: [[0.1,0.4]] + """ + if not get_config_context().validation_enabled: + return check_obj + + with config_context(validation_depth=get_validation_depth(check_obj)): + output = self.get_backend(check_obj).validate( + check_obj=check_obj, + schema=self, + head=head, + tail=tail, + sample=sample, + random_state=random_state, + lazy=lazy, + inplace=inplace, + ) + + return output + + @_DataFrameSchema.dtype.setter # type: ignore[attr-defined] + def dtype(self, value) -> None: + """Set the dtype property.""" + self._dtype = resolve_dtype(value) + + def strategy(self, *, size: int | None = None, n_regex_columns: int = 1): + """Data synthesis is not supported for pyarrow schemas.""" + raise NotImplementedError( + "Data synthesis is not supported with pyarrow schemas." + ) + + def example(self, size: int | None = None, n_regex_columns: int = 1): + """Data synthesis is not supported for pyarrow schemas.""" + raise NotImplementedError( + "Data synthesis is not supported with pyarrow schemas." + ) diff --git a/pandera/api/pyarrow/model.py b/pandera/api/pyarrow/model.py new file mode 100644 index 000000000..4f6ce1d81 --- /dev/null +++ b/pandera/api/pyarrow/model.py @@ -0,0 +1,160 @@ +"""Class-based API for PyArrow models.""" + +from __future__ import annotations + +import inspect +import sys +from typing import cast + +import pyarrow as pa + +from pandera.api.base.schema import BaseSchema +from pandera.api.checks import Check +from pandera.api.dataframe.model import DataFrameModel as _DataFrameModel +from pandera.api.dataframe.model import _dtype_metadata, get_dtype_kwargs +from pandera.api.dataframe.model_components import FieldInfo +from pandera.api.pyarrow.components import Column +from pandera.api.pyarrow.container import DataFrameSchema +from pandera.api.pyarrow.model_config import BaseConfig +from pandera.api.pyarrow.utils import resolve_dtype +from pandera.engines import narwhals_engine +from pandera.errors import SchemaInitError +from pandera.typing import AnnotationInfo +from pandera.typing.pyarrow import Table +from pandera.utils import docstring_substitution + +if sys.version_info < (3, 11): + from typing_extensions import Self +else: + from typing import Self + + +class DataFrameModel(_DataFrameModel[pa.Table, DataFrameSchema]): + """Definition of a :class:`~pandera.api.pyarrow.container.DataFrameSchema`. + + See the :ref:`User Guide ` for more. + """ + + Config: type[BaseConfig] = BaseConfig + + @classmethod + def build_schema_(cls, **kwargs) -> DataFrameSchema: + return DataFrameSchema( + cls._build_columns(cls.__fields__, cls.__checks__), + checks=cls.__root_checks__, + **kwargs, + ) + + @classmethod + def _build_columns( + cls, + fields: dict[str, tuple[AnnotationInfo, FieldInfo]], + checks: dict[str, list[Check]], + ) -> dict[str, Column]: + columns: dict[str, Column] = {} + for field_name, (annotation, field) in fields.items(): + field_checks = checks.get(field_name, []) + field_name = field.name + check_name = getattr(field, "check_name", None) + + try: + engine_dtype = resolve_dtype(annotation.raw_annotation) + if inspect.isclass(annotation.raw_annotation) and issubclass( + annotation.raw_annotation, narwhals_engine.DataType + ): + # use the raw annotation as the dtype if it's a native + # pandera narwhals datatype + dtype = annotation.raw_annotation + else: + dtype = engine_dtype.type + except (TypeError, ValueError) as exc: + if annotation.metadata: + if field.dtype_kwargs: + raise TypeError( + "Cannot specify redundant 'dtype_kwargs' " + + f"for {annotation.raw_annotation}." + + "\n Usage Tip: Drop 'typing.Annotated'." + ) from exc + # ``Annotated`` may carry only a FieldInfo (e.g. + # ``Annotated[float, pa.Field(...)]``) without any + # dtype parameters. In that case, use the annotated + # type as-is. + if _dtype_metadata(annotation): + dtype_kwargs = get_dtype_kwargs(annotation) + dtype = annotation.arg(**dtype_kwargs) # type: ignore + else: + dtype = annotation.arg # type: ignore + elif annotation.default_dtype: + dtype = annotation.default_dtype + else: + dtype = annotation.arg # type: ignore + + if annotation.origin is None or dtype: + if check_name is False: + raise SchemaInitError( + f"'check_name' is not supported for {field_name}." + ) + + column_kwargs = ( + field.column_properties( + dtype, + required=not annotation.optional, + checks=field_checks, + name=field_name, + ) + if field + else {} + ) + columns[field_name] = Column(**column_kwargs) + + else: + raise SchemaInitError( + f"Invalid annotation '{field_name}: " + f"{annotation.raw_annotation}'." + ) + + return columns + + @classmethod + @docstring_substitution(validate_doc=BaseSchema.validate.__doc__) + def validate( + cls: type[Self], + check_obj: pa.Table, + head: int | None = None, + tail: int | None = None, + sample: int | None = None, + random_state: int | None = None, + lazy: bool = False, + inplace: bool = False, + ) -> Table[Self]: + """%(validate_doc)s""" + result = cls.to_schema().validate( + check_obj, head, tail, sample, random_state, lazy, inplace + ) + return cast(Table[Self], result) + + @classmethod + def empty(cls: type[Self], *_args) -> Table[Self]: + """Create an empty pyarrow Table with the schema of this model.""" + schema = cls.to_schema() + arrow_schema = pa.schema( + [ + (name, _narwhals_dtype_to_pyarrow(col.dtype)) + for name, col in schema.columns.items() + ] + ) + return cast(Table[Self], arrow_schema.empty_table()) + + +def _narwhals_dtype_to_pyarrow(dtype) -> pa.DataType: + """Translate a narwhals engine dtype back into a ``pyarrow.DataType``. + + Narwhals owns the pyarrow mapping, so build a zero-row frame of the + narwhals dtype and read the pyarrow type back off it rather than + maintaining a parallel lookup table here. + """ + import narwhals.stable.v1 as nw + + empty = nw.from_native(pa.table({"x": pa.array([])}), eager_only=True) + casted = empty.with_columns(nw.col("x").cast(dtype.type)) + return nw.to_native(casted).schema.field("x").type diff --git a/pandera/api/pyarrow/model_config.py b/pandera/api/pyarrow/model_config.py new file mode 100644 index 000000000..909fc19a6 --- /dev/null +++ b/pandera/api/pyarrow/model_config.py @@ -0,0 +1,14 @@ +"""Class-based dataframe model API configuration for pyarrow.""" + +from __future__ import annotations + +from pandera.api.dataframe.model_config import BaseConfig as _BaseConfig +from pandera.api.pyarrow.types import PyArrowDtypeInputTypes + + +class BaseConfig(_BaseConfig): + """Define pyarrow DataFrameSchema-wide options.""" + + #: datatype of the table. This overrides the data types specified in + #: any of the fields. + dtype: PyArrowDtypeInputTypes | None = None diff --git a/pandera/api/pyarrow/types.py b/pandera/api/pyarrow/types.py new file mode 100644 index 000000000..e47bbfdcf --- /dev/null +++ b/pandera/api/pyarrow/types.py @@ -0,0 +1,35 @@ +"""PyArrow types.""" + +from typing import NamedTuple, Union + +import pyarrow as pa + + +class PyArrowData(NamedTuple): + """Data container passed to ``native=True`` pyarrow checks. + + Mirrors :class:`~pandera.api.polars.types.PolarsData` and + :class:`~pandera.api.ibis.types.IbisData` so that check functions taking a + single positional argument receive the same shape across backends. + """ + + table: pa.Table + key: str = "*" + + +class CheckResult(NamedTuple): + """Check result for user-defined checks.""" + + check_output: pa.Table + check_passed: pa.Table + checked_object: pa.Table + failure_cases: pa.Table + + +PyArrowCheckObjects = pa.Table + +PyArrowDtypeInputTypes = Union[ + str, + type, + pa.DataType, +] diff --git a/pandera/api/pyarrow/utils.py b/pandera/api/pyarrow/utils.py new file mode 100644 index 000000000..1545ebf52 --- /dev/null +++ b/pandera/api/pyarrow/utils.py @@ -0,0 +1,63 @@ +"""Utilities for the pyarrow schema API.""" + +from __future__ import annotations + +from typing import Any + +import narwhals.stable.v1 as nw +import pyarrow as pa + +from pandera.config import ( + ValidationDepth, + get_config_context, + get_config_global, +) +from pandera.engines import narwhals_engine + + +def pyarrow_dtype_to_narwhals(dtype: pa.DataType) -> Any: + """Translate a native ``pyarrow.DataType`` into its narwhals equivalent. + + ``narwhals_engine`` resolves strings, narwhals dtypes and abstract pandera + dtypes, but only a handful of pyarrow types happen to stringify into + something it recognises: ``pa.int64()`` renders as ``"int64"`` and + resolves, while ``pa.float64()`` renders as ``"double"`` and + ``pa.date32()`` as ``"date32[day]"``, neither of which do. Round-tripping a + zero-length array through narwhals yields an exact translation for every + pyarrow type, including parametrized ones (timestamps, decimals, lists, + structs). + """ + empty = pa.table({"x": pa.array([], type=dtype)}) + return nw.from_native(empty, eager_only=True).schema["x"] + + +def resolve_dtype(value: Any): + """Resolve a user-supplied dtype into a narwhals engine ``DataType``. + + Accepts anything ``narwhals_engine`` understands plus native pyarrow + types, which are translated first. + """ + if value is None: + return None + if isinstance(value, pa.DataType): + value = pyarrow_dtype_to_narwhals(value) + return narwhals_engine.Engine.dtype(value) + + +def get_validation_depth(check_obj: pa.Table) -> ValidationDepth: + """Get the validation depth for a pyarrow Table. + + Unlike ``pl.LazyFrame`` or ``ibis.Table``, a ``pyarrow.Table`` is always + fully materialized in memory, so running data-level checks by default + costs nothing extra — it gets the same treatment as ``pl.DataFrame``. + Explicit context or global configuration still wins. + """ + config_ctx = get_config_context(validation_depth_default=None) + if config_ctx.validation_depth is not None: + return config_ctx.validation_depth + + config_global = get_config_global() + if config_global.validation_depth is not None: + return config_global.validation_depth + + return ValidationDepth.SCHEMA_AND_DATA diff --git a/pandera/backends/narwhals/base.py b/pandera/backends/narwhals/base.py index 0b8bce082..990de5ac2 100644 --- a/pandera/backends/narwhals/base.py +++ b/pandera/backends/narwhals/base.py @@ -33,6 +33,32 @@ except ImportError: # pragma: no cover — polars is optional pl = None # type: ignore[assignment] +try: + import pyarrow as _pa # used in the pyarrow failure_cases paths +except ImportError: # pragma: no cover — pyarrow is optional + _pa = None # type: ignore[assignment] + + +def _errors_implementation(schema_errors: list[SchemaError]) -> Any: + """Infer the frame backend a batch of errors came from. + + Scalar failure cases (missing column, wrong dtype, …) carry no frame, so + their builder has nothing to dispatch on. Taking the implementation from + whichever errors *do* carry a frame keeps every piece of one batch on the + same backend, which is what ``_concat_failure_cases`` needs to combine + them. Returns ``None`` when no error in the batch carries a frame. + """ + for err in schema_errors: + try: + fc = nw.from_native( + err.failure_cases, eager_or_interchange_only=False + ) + except TypeError: + continue + if isinstance(fc, (nw.DataFrame, nw.LazyFrame)): + return fc.implementation + return None + def _lit_nullable_int32(value, implementation) -> Any: """Int32 literal that tolerates ``None`` on pandas-like backends. @@ -148,6 +174,12 @@ def _concat_failure_cases(items: list, implementation=None) -> Any: ) native_items = [nw.to_native(item) for item in nw_items] return functools.reduce(lambda a, b: a.union(b), native_items) + elif first_nw.implementation == nw.Implementation.PYARROW: + # pyarrow.Table has no ``.union()``; concatenate via narwhals and + # hand back the native table. Every piece of a pyarrow run is + # built by the pyarrow builders above, so there are no pl_items + # to merge here. + return nw.to_native(nw.concat(nw_items)) elif first_nw.implementation == nw.Implementation.POLARS: # Polars lazy path: use nw.concat to stay lazy, then unwrap. # When pl_items are also present (schema-level failure cases from @@ -418,6 +450,7 @@ def failure_cases_metadata( conversion, no Arrow roundtrip for lazy/SQL backends. """ failure_case_collection: list = [] + errors_implementation = _errors_implementation(schema_errors) # Implementation of the validated frame — used to disambiguate the # all-scalar failure-cases path in _concat_failure_cases (e.g. pandas @@ -472,8 +505,34 @@ def failure_cases_metadata( self._build_lazy_failure_case(fc, err, check_identifier) ) elif isinstance(fc, (nw.LazyFrame, nw.DataFrame)): + if fc.implementation == nw.Implementation.PYARROW: + failure_case_collection.append( + self._build_pyarrow_failure_case( + fc, err, check_identifier + ) + ) + else: + failure_case_collection.append( + self._build_eager_failure_case( + fc, err, check_identifier + ) + ) + elif errors_implementation == nw.Implementation.PYARROW or ( + pl is None + and _pa is not None + # pandas-like runs keep _build_scalar_failure_case, which + # already builds a pandas frame when polars is absent. + # ``errors_implementation`` is None for an all-scalar batch, + # so the validated frame's implementation is the only usable + # signal here; without it a pandas + pyarrow (no polars) + # install would report failure_cases as a pyarrow.Table. + and data_implementation + not in _EAGER_PANDAS_LIKE_IMPLEMENTATIONS + ): failure_case_collection.append( - self._build_eager_failure_case(fc, err, check_identifier) + self._build_pyarrow_scalar_failure_case( + err, check_identifier + ) ) else: failure_case_collection.append( @@ -548,6 +607,127 @@ def _build_lazy_failure_case(fc, err: SchemaError, check_identifier): # module-string sniffing. return enriched + @staticmethod + def _resolved_check_output(err: SchemaError): + """Normalize ``err.check_output`` to a Narwhals frame, or ``None``.""" + if err.check_output is None: + return None + co = err.check_output + if isinstance(co, (nw.LazyFrame, nw.DataFrame)): + return co + if isinstance(co, nw.Expr): + return None + return nw.from_native(co, eager_or_interchange_only=False) + + @staticmethod + def _failing_row_indices(err: SchemaError) -> list | None: + """Positions of the rows that failed the check, if recoverable.""" + resolved_co = NarwhalsSchemaBackend._resolved_check_output(err) + if resolved_co is None: + return None + co_eager = _materialize(resolved_co) + try: + co_indexed = co_eager.with_row_index("index") + except AttributeError: + # Older polars: ``with_row_index`` was called ``with_row_count``. + co_indexed = co_eager.with_row_count("index") + return co_indexed.filter(~nw.col(CHECK_OUTPUT_KEY))["index"].to_list() + + @staticmethod + def _build_pyarrow_failure_case(fc, err: SchemaError, check_identifier): + """Build an eager pyarrow failure-case table with row-index enrichment. + + Mirrors ``_build_eager_failure_case`` but stays in pyarrow rather than + round-tripping through polars, so that a ``pandera[pyarrow]`` install + without polars can still report failure cases — and so the reported + frame does not change type depending on whether polars happens to be + installed. + + Returns a Narwhals-wrapped frame so ``_concat_failure_cases`` can + dispatch on ``item.implementation``. + """ + import json + + fc_native = nw.to_native(_materialize(fc)) + rows = fc_native.to_pylist() + col_names = list(fc_native.column_names) + + if len(col_names) == 1: + key = col_names[0] + failure_case = [ + None if row[key] is None else str(row[key]) for row in rows + ] + else: + # Match the polars path, which JSON-encodes the whole row when a + # failure case spans multiple columns. + failure_case = [json.dumps(row, default=str) for row in rows] + + failing = NarwhalsSchemaBackend._failing_row_indices(err) + if failing is None: + index: list = [None] * len(rows) + else: + index = list(failing[: len(rows)]) + index += [None] * (len(rows) - len(index)) + + table = _pa.table( + { + "failure_case": _pa.array(failure_case, type=_pa.string()), + "schema_context": _pa.array( + [err.schema.__class__.__name__] * len(rows), + type=_pa.string(), + ), + "column": _pa.array( + [err.schema.name] * len(rows), type=_pa.string() + ), + "check": _pa.array( + [ + None + if check_identifier is None + else str(check_identifier) + ] + * len(rows), + type=_pa.string(), + ), + "check_number": _pa.array( + [err.check_index] * len(rows), type=_pa.int32() + ), + "index": _pa.array(index, type=_pa.int32()), + } + ) + return nw.from_native(table, eager_only=True) + + @staticmethod + def _build_pyarrow_scalar_failure_case(err: SchemaError, check_identifier): + """Scalar failure case as a one-row pyarrow table. + + The polars equivalent is ``_build_scalar_failure_case``; this keeps a + pyarrow validation run on a single backend so the pieces concatenate. + """ + failure_case = err.failure_cases + table = _pa.table( + { + "failure_case": _pa.array( + [None if failure_case is None else str(failure_case)], + type=_pa.string(), + ), + "schema_context": _pa.array( + [err.schema.__class__.__name__], type=_pa.string() + ), + "column": _pa.array([err.schema.name], type=_pa.string()), + "check": _pa.array( + [ + None + if check_identifier is None + else str(check_identifier) + ], + type=_pa.string(), + ), + "check_number": _pa.array([err.check_index], type=_pa.int32()), + "index": _pa.array([None], type=_pa.int32()), + } + ) + return nw.from_native(table, eager_only=True) + @staticmethod def _build_eager_failure_case(fc, err: SchemaError, check_identifier): """Build an eager polars failure-case frame with row-index enrichment. diff --git a/pandera/backends/narwhals/checks.py b/pandera/backends/narwhals/checks.py index 90e4661c7..86e701fdd 100644 --- a/pandera/backends/narwhals/checks.py +++ b/pandera/backends/narwhals/checks.py @@ -107,6 +107,16 @@ def _is_ibis_native(frame: Any) -> bool: return isinstance(frame, ibis.Table) +def _is_pyarrow_native(frame: Any) -> bool: + """Cheap ``pyarrow.Table`` detection that mirrors ``_is_polars_native``. + + Duck-types via the module name so the narwhals backend keeps working on + installs where pyarrow is not present. + """ + mod = getattr(type(frame), "__module__", "") or "" + return mod.startswith("pyarrow") and type(frame).__name__ == "Table" + + # Root modules of pandas-like libraries. Compared against the first dotted # segment of ``type(obj).__module__``: pandas < 3 reports e.g. # "pandas.core.frame" while pandas >= 3 reports just "pandas" for its public @@ -158,6 +168,11 @@ def _wrap_native_frame_with_key(native_frame: Any, key: str | None) -> Any: return IbisData(table=native_frame, key=key) + if _is_pyarrow_native(native_frame): + from pandera.api.pyarrow.types import PyArrowData + + return PyArrowData(table=native_frame, key=key or "*") + if _is_pandas_like_native(native_frame): # pandas-style user check functions receive exactly what the native # pandas backend passes them: the Series for a column-level check and @@ -439,6 +454,31 @@ def _normalize_native_output(out, check_obj: NarwhalsData): native.with_columns(bool_col), eager_only=True ) + # Handle pyarrow native return types from native=True checks. As with + # polars above, pyarrow is not imported at module level so ibis-only / + # polars-only installs keep working — detection is by module name. + if out_mod.startswith("pyarrow"): + out_type_name = type(out).__name__ + + if out_type_name == "Table": + return nw.from_native(out, eager_only=True) + + if out_type_name.endswith("Scalar"): + # Aggregate boolean result — hand back a plain Python bool. + return bool(out.as_py()) + + if out_type_name.endswith(("Array", "ChunkedArray")): + # Row-level boolean output — attach it to the original table as + # ``CHECK_OUTPUT_KEY`` to build a wide frame, mirroring the + # polars branch above. + from pandera.api.narwhals.utils import _materialize + + native = nw.to_native(_materialize(check_obj.frame)) + tbl = native.append_column(CHECK_OUTPUT_KEY, out) + return nw.from_native(tbl, eager_only=True) + + return out # pragma: no cover — unexpected pyarrow type + # Handle pandas-like native return types (pandas, modin, cuDF) from # native=True checks: boolean Series / DataFrame outputs are attached # to the original frame as CHECK_OUTPUT_KEY, mirroring the polars diff --git a/pandera/backends/pyarrow/__init__.py b/pandera/backends/pyarrow/__init__.py new file mode 100644 index 000000000..9a6aa2c81 --- /dev/null +++ b/pandera/backends/pyarrow/__init__.py @@ -0,0 +1,7 @@ +"""PyArrow backend implementation for schemas and checks. + +Validation itself is performed by the narwhals backends +(:mod:`pandera.backends.narwhals`), which are frame-agnostic; this package +only wires ``pyarrow.Table`` into the backend registry, lazily via +:func:`~pandera.backends.pyarrow.register.register_pyarrow_backends`. +""" diff --git a/pandera/backends/pyarrow/register.py b/pandera/backends/pyarrow/register.py new file mode 100644 index 000000000..c8640f606 --- /dev/null +++ b/pandera/backends/pyarrow/register.py @@ -0,0 +1,44 @@ +"""Register PyArrow backends.""" + +from __future__ import annotations + +from functools import lru_cache + +import pyarrow as pa + + +@lru_cache +def register_pyarrow_backends(check_cls_fqn: str | None = None): + """Register backends for ``pyarrow.Table``. + + Unlike polars and ibis — which have hand-written native backends and only + use narwhals when ``PANDERA_USE_NARWHALS_BACKEND=True`` — pyarrow is served + exclusively by the narwhals backends, so narwhals is a hard requirement of + the pyarrow API. + + Decorated with ``@lru_cache`` to prevent duplicate registrations across + repeated ``validate()`` calls. + """ + try: + import narwhals.stable.v1 as nw + except ImportError as exc: # pragma: no cover — narwhals is a dependency + raise ImportError( + "The pyarrow schema API requires the 'narwhals' package. " + "Install it with: pip install 'pandera[pyarrow]'" + ) from exc + + import pandera.backends.narwhals.builtin_checks # noqa: F401 + from pandera.api.checks import Check + from pandera.api.pyarrow.components import Column + from pandera.api.pyarrow.container import DataFrameSchema + from pandera.backends.narwhals.checks import NarwhalsCheckBackend + from pandera.backends.narwhals.components import ColumnBackend + from pandera.backends.narwhals.container import DataFrameSchemaBackend + + DataFrameSchema.register_backend( + pa.Table, DataFrameSchemaBackend, force=True + ) + Column.register_backend(pa.Table, ColumnBackend, force=True) + Check.register_backend(pa.Table, NarwhalsCheckBackend, force=True) + Check.register_backend(nw.LazyFrame, NarwhalsCheckBackend, force=True) + Check.register_backend(nw.DataFrame, NarwhalsCheckBackend, force=True) diff --git a/pandera/backends/register_checks.py b/pandera/backends/register_checks.py index 36423ba2c..3fabc94f4 100644 --- a/pandera/backends/register_checks.py +++ b/pandera/backends/register_checks.py @@ -40,9 +40,28 @@ def register_default_check_backends(check_obj_cls: type) -> None: # Narwhals wrapper types share pandas-like class names (e.g. DataFrame) # but must not route through the pandas backend registry. if module.startswith("narwhals."): + from pandera.api.checks import Check from pandera.config import CONFIG + # Fast path: a check backend for this narwhals wrapper type is already + # registered, which is the only thing ``Check.get_backend`` looks up. + # Returning here keeps a check on a narwhals frame from eagerly pulling + # *every* narwhals-compatible library's backends into the shared + # registry — that side effect primes each library's registration + # ``lru_cache``, so a later wipe of ``BACKEND_REGISTRY`` leaves those + # libraries permanently unregistered. + if (Check, check_obj_cls) in Check.BACKEND_REGISTRY: + return + use_nw = CONFIG.use_narwhals_backend + try: + from pandera.backends.pyarrow.register import ( + register_pyarrow_backends, + ) + + register_pyarrow_backends() + except ImportError: + pass try: from pandera.backends.polars.register import ( register_polars_backends, @@ -93,6 +112,14 @@ def register_default_check_backends(check_obj_cls: type) -> None: ) return + if module.startswith("pyarrow"): + from pandera.backends.pyarrow.register import ( + register_pyarrow_backends, + ) + + register_pyarrow_backends() + return + if module.startswith("ibis."): from pandera.backends.ibis.register import register_ibis_backends from pandera.config import CONFIG diff --git a/pandera/engines/narwhals_engine.py b/pandera/engines/narwhals_engine.py index 95107d292..00dbadb0a 100644 --- a/pandera/engines/narwhals_engine.py +++ b/pandera/engines/narwhals_engine.py @@ -167,7 +167,7 @@ class Int32(DataType, dtypes.Int32): @Engine.register_dtype( - equivalents=["int64", nw.Int64, dtypes.Int64, dtypes.Int64()] + equivalents=["int64", int, nw.Int64, dtypes.Int64, dtypes.Int64()] ) @immutable class Int64(DataType, dtypes.Int64): @@ -232,7 +232,13 @@ class Float32(DataType, dtypes.Float32): @Engine.register_dtype( - equivalents=["float64", nw.Float64, dtypes.Float64, dtypes.Float64()] + equivalents=[ + "float64", + float, + nw.Float64, + dtypes.Float64, + dtypes.Float64(), + ] ) @immutable class Float64(DataType, dtypes.Float64): @@ -247,7 +253,14 @@ class Float64(DataType, dtypes.Float64): @Engine.register_dtype( - equivalents=["str", "string", nw.String, dtypes.String, dtypes.String()] + equivalents=[ + "str", + "string", + str, + nw.String, + dtypes.String, + dtypes.String(), + ] ) @immutable class String(DataType, dtypes.String): @@ -257,7 +270,7 @@ class String(DataType, dtypes.String): @Engine.register_dtype( - equivalents=["bool", nw.Boolean, dtypes.Bool, dtypes.Bool()] + equivalents=["bool", bool, nw.Boolean, dtypes.Bool, dtypes.Bool()] ) @immutable class Bool(DataType, dtypes.Bool): diff --git a/pandera/pyarrow.py b/pandera/pyarrow.py new file mode 100644 index 000000000..954fce44a --- /dev/null +++ b/pandera/pyarrow.py @@ -0,0 +1,35 @@ +"""A flexible and expressive pyarrow.Table validation library.""" + +import pandera.backends.pyarrow +from pandera import errors +from pandera.api.checks import Check +from pandera.api.dataframe.model_components import ( + Field, + check, + dataframe_check, +) +from pandera.api.pyarrow.components import Column +from pandera.api.pyarrow.container import DataFrameSchema +from pandera.api.pyarrow.model import DataFrameModel +from pandera.api.pyarrow.types import PyArrowData +from pandera.config import set_config +from pandera.decorators import check_input, check_io, check_output, check_types +from pandera.typing import pyarrow as typing + +__all__ = [ + "check_input", + "check_io", + "check_output", + "check_types", + "check", + "Check", + "Column", + "dataframe_check", + "DataFrameModel", + "DataFrameSchema", + "errors", + "Field", + "PyArrowData", + "set_config", + "typing", +] diff --git a/pandera/typing/pyarrow.py b/pandera/typing/pyarrow.py new file mode 100644 index 000000000..d2f1ed589 --- /dev/null +++ b/pandera/typing/pyarrow.py @@ -0,0 +1,162 @@ +"""Pandera type annotations for PyArrow.""" + +import functools +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +from packaging import version + +from pandera.typing.common import DataFrameBase, DataFrameModel +from pandera.typing.formats import Formats + +try: + import pyarrow + + PYARROW_INSTALLED = True +except ImportError: + PYARROW_INSTALLED = False + + +def pyarrow_version(): + """Return the pyarrow version.""" + return version.parse(pyarrow.__version__) + + +if TYPE_CHECKING: + T = TypeVar("T") # pragma: no cover +else: + T = DataFrameModel + + +if PYARROW_INSTALLED: + + class Table(DataFrameBase, pyarrow.Table, Generic[T]): + """Pandera generic for ``pyarrow.Table``, only used for annotation.""" + + @classmethod + def from_format(cls, obj: Any, config) -> "pyarrow.Table": + """Convert serialized data into a ``pyarrow.Table``. + + The format is taken from the + :py:class:`pandera.api.pyarrow.model.DataFrameModel` config + options ``from_format`` and ``from_format_kwargs``. + + :param obj: object representing a serialized table. + :param config: dataframe model configuration object. + """ + if config.from_format is None: + if not isinstance(obj, pyarrow.Table): + try: + obj = pyarrow.table(obj) + except Exception as exc: + raise ValueError( + f"Expected pyarrow.Table, found {type(obj)}" + ) from exc + return obj + + if callable(config.from_format): + reader = config.from_format + return reader(obj, **(config.from_format_kwargs or {})) + + try: + format_type = Formats(config.from_format) + except ValueError as exc: + raise ValueError( + f"Unsupported format: {config.from_format}. " + "PyArrow natively supports: dict, csv, json, and parquet." + ) from exc + + kwargs = config.from_format_kwargs or {} + + if format_type == Formats.dict: + if not isinstance(obj, dict): + raise ValueError( + f"Expected dict for dict format, got {type(obj)}" + ) + return pyarrow.table(obj) + + try: + if format_type == Formats.csv: + from pyarrow import csv + + return csv.read_csv(obj, **kwargs) + if format_type == Formats.json: + from pyarrow import json as pa_json + + return pa_json.read_json(obj, **kwargs) + if format_type == Formats.parquet: + from pyarrow import parquet + + return parquet.read_table(obj, **kwargs) + if format_type == Formats.feather: + from pyarrow import feather + + return feather.read_table(obj, **kwargs) + except Exception as exc: + raise ValueError( + f"Failed to read {format_type.value} with PyArrow: {exc}" + ) from exc + + raise ValueError( + f"{format_type.value} format is not natively supported by " + "PyArrow. Use a custom callable for from_format instead." + ) + + @classmethod + def to_format(cls, data: "pyarrow.Table", config) -> Any: + """Convert a table to the format specified in the model config. + + Driven by the + :py:class:`pandera.api.pyarrow.model.DataFrameModel` config + options ``to_format`` and ``to_format_kwargs``. + + :param data: convert this data to the specified format + :param config: config object from the DataFrameModel + """ + if config.to_format is None: + return data + + if callable(config.to_format): + writer = functools.partial(config.to_format, data) + buffer = ( + config.to_format_buffer() + if callable(config.to_format_buffer) + else None + ) + args = [] if buffer is None else [buffer] + out = writer(*args, **(config.to_format_kwargs or {})) + return out if buffer is None else buffer + + try: + format_type = Formats(config.to_format) + except ValueError as exc: + raise ValueError( + f"Unsupported format: {config.to_format}. " + "PyArrow natively supports: dict, parquet, and feather." + ) from exc + + kwargs = config.to_format_kwargs or {} + + if format_type == Formats.dict: + return data.to_pydict() + if format_type == Formats.parquet: + from pyarrow import parquet + + return parquet.write_table(data, **kwargs) + if format_type == Formats.feather: + from pyarrow import feather + + return feather.write_feather(data, **kwargs) + + raise ValueError( + f"{format_type.value} format is not natively supported by " + "PyArrow. Use a custom callable for to_format instead." + ) + + @classmethod + def _get_schema_model(cls, field): + if not field.sub_fields: + raise TypeError( + "Expected a typed pandera.typing.pyarrow.Table," + " e.g. Table[Schema]" + ) + return field.sub_fields[0].type_ diff --git a/pyproject.toml b/pyproject.toml index 07436c9e5..405b7bbf9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,10 @@ xarray = [ "numpy >= 1.24.4", ] narwhals = ["narwhals >= 1.26.0"] +pyarrow = [ + "pyarrow >= 13", + "narwhals >= 1.26.0", +] all = [ "hypothesis >= 6.92.7", "scipy", @@ -121,6 +125,7 @@ all = [ "xarray >= 2024.10.0", "numpy >= 1.24.4", "narwhals >= 1.26.0", + "pyarrow >= 13", ] [dependency-groups] diff --git a/tests/pyarrow/__init__.py b/tests/pyarrow/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/pyarrow/test_pyarrow_checks.py b/tests/pyarrow/test_pyarrow_checks.py new file mode 100644 index 000000000..444136bfe --- /dev/null +++ b/tests/pyarrow/test_pyarrow_checks.py @@ -0,0 +1,124 @@ +"""Tests for checks against pyarrow Tables.""" + +import pyarrow +import pyarrow.compute as pc +import pytest + +import pandera.pyarrow as pa +from pandera.api.pyarrow.types import PyArrowData +from pandera.errors import SchemaError + + +@pytest.mark.parametrize( + "check,passing,failing", + [ + (pa.Check.gt(0), [1, 2], [0, 1]), + (pa.Check.ge(0), [0, 1], [-1, 0]), + (pa.Check.lt(10), [1, 2], [10, 1]), + (pa.Check.le(10), [10], [11]), + (pa.Check.eq(1), [1, 1], [1, 2]), + (pa.Check.ne(1), [2, 3], [1, 2]), + (pa.Check.isin([1, 2]), [1, 2], [1, 3]), + (pa.Check.notin([3]), [1, 2], [3]), + (pa.Check.between(0, 5), [1, 5], [6]), + ], +) +def test_builtin_numeric_checks(check, passing, failing): + schema = pa.DataFrameSchema({"a": pa.Column(int, check)}) + assert schema.validate(pyarrow.table({"a": passing})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": failing})) + + +@pytest.mark.parametrize( + "check,passing,failing", + [ + (pa.Check.str_startswith("a"), ["ab", "ac"], ["ab", "bc"]), + (pa.Check.str_endswith("z"), ["az"], ["za"]), + (pa.Check.str_contains("b"), ["abc"], ["acd"]), + (pa.Check.str_matches(r"^\d+$"), ["123"], ["12a"]), + ], +) +def test_builtin_string_checks(check, passing, failing): + schema = pa.DataFrameSchema({"a": pa.Column(str, check)}) + assert schema.validate(pyarrow.table({"a": passing})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": failing})) + + +def test_native_check_pyarrow_data_convention(): + """A 1-arg native check receives a ``PyArrowData`` container. + + This mirrors ``PolarsData`` / ``IbisData`` on the other backends. + """ + seen = {} + + def check_fn(data): + seen["type"] = type(data) + seen["key"] = data.key + return pc.greater(data.table[data.key], 0) + + schema = pa.DataFrameSchema({"a": pa.Column(int, pa.Check(check_fn))}) + schema.validate(pyarrow.table({"a": [1, 2]})) + + assert seen["type"] is PyArrowData + assert seen["key"] == "a" + + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, -2]})) + + +def test_native_check_two_arg_convention(): + """A 2-arg native check gets ``(native_table, key)``.""" + schema = pa.DataFrameSchema( + { + "a": pa.Column( + int, pa.Check(lambda tbl, key: pc.greater(tbl[key], 0)) + ) + } + ) + assert schema.validate(pyarrow.table({"a": [1, 2]})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, -2]})) + + +def test_native_check_returning_boolean_scalar(): + """An aggregate pyarrow scalar is normalized to a python bool.""" + schema = pa.DataFrameSchema( + { + "a": pa.Column( + int, + pa.Check(lambda tbl, key: pc.all(pc.greater(tbl[key], 0))), + ) + } + ) + assert schema.validate(pyarrow.table({"a": [1, 2]})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, -2]})) + + +def test_non_native_expression_check(): + """``native=False`` uses the narwhals expression protocol.""" + schema = pa.DataFrameSchema( + {"a": pa.Column(int, pa.Check(lambda col: col > 0, native=False))} + ) + assert schema.validate(pyarrow.table({"a": [1, 2]})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, -2]})) + + +def test_element_wise_check(): + schema = pa.DataFrameSchema( + {"a": pa.Column(int, pa.Check(lambda v: v > 0, element_wise=True))} + ) + assert schema.validate(pyarrow.table({"a": [1, 2]})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, -2]})) + + +def test_failure_cases_are_reported(): + schema = pa.DataFrameSchema({"a": pa.Column(int, pa.Check.gt(0))}) + with pytest.raises(SchemaError) as exc_info: + schema.validate(pyarrow.table({"a": [1, -2, -3]}), lazy=False) + failure_cases = exc_info.value.failure_cases + assert failure_cases is not None diff --git a/tests/pyarrow/test_pyarrow_container.py b/tests/pyarrow/test_pyarrow_container.py new file mode 100644 index 000000000..d68901a56 --- /dev/null +++ b/tests/pyarrow/test_pyarrow_container.py @@ -0,0 +1,252 @@ +"""Tests for the pyarrow DataFrameSchema API.""" + +import pyarrow +import pyarrow.compute as pc +import pytest + +import pandera.pyarrow as pa +from pandera.config import ValidationDepth, config_context +from pandera.errors import SchemaError, SchemaErrors, SchemaWarning + + +@pytest.fixture +def table(): + return pyarrow.table( + { + "int_col": [1, 2, 3], + "float_col": [1.0, 2.0, 3.0], + "str_col": ["a", "b", "c"], + } + ) + + +@pytest.fixture +def schema(): + return pa.DataFrameSchema( + { + "int_col": pa.Column(int, pa.Check.gt(0)), + "float_col": pa.Column(float), + "str_col": pa.Column(str), + } + ) + + +def test_validate_returns_pyarrow_table(table, schema): + validated = schema.validate(table) + assert isinstance(validated, pyarrow.Table) + assert validated.equals(table) + + +def test_data_level_check_failure_is_raised(schema): + """Regression test: data-level checks must not be silently skipped. + + A pyarrow.Table is always materialized, so the default validation depth + is SCHEMA_AND_DATA — the same as pl.DataFrame, not pl.LazyFrame. + """ + invalid = pyarrow.table( + { + "int_col": [1, -2, 3], + "float_col": [1.0, 2.0, 3.0], + "str_col": ["a", "b", "c"], + } + ) + with pytest.raises(SchemaError, match="greater_than"): + schema.validate(invalid) + + +def test_default_validation_depth_is_schema_and_data(table): + from pandera.api.pyarrow.utils import get_validation_depth + + assert get_validation_depth(table) is ValidationDepth.SCHEMA_AND_DATA + + +def test_validation_depth_config_is_respected(schema): + invalid = pyarrow.table( + { + "int_col": [1, -2, 3], + "float_col": [1.0, 2.0, 3.0], + "str_col": ["a", "b", "c"], + } + ) + with config_context(validation_depth=ValidationDepth.SCHEMA_ONLY): + assert schema.validate(invalid).equals(invalid) + + +def test_wrong_dtype_raises(schema): + invalid = pyarrow.table( + { + "int_col": [1.5, 2.5, 3.5], + "float_col": [1.0, 2.0, 3.0], + "str_col": ["a", "b", "c"], + } + ) + with pytest.raises(SchemaError): + schema.validate(invalid) + + +def test_missing_column_raises(schema): + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"int_col": [1, 2, 3]})) + + +def test_strict_rejects_extra_column(table): + schema = pa.DataFrameSchema({"int_col": pa.Column(int)}, strict=True) + with pytest.raises(SchemaError): + schema.validate(table) + + +def test_strict_filter_drops_extra_columns(table): + schema = pa.DataFrameSchema({"int_col": pa.Column(int)}, strict="filter") + validated = schema.validate(table) + assert validated.column_names == ["int_col"] + + +def test_lazy_collects_all_errors(): + schema = pa.DataFrameSchema( + { + "a": pa.Column(int, pa.Check.gt(10)), + "b": pa.Column(str), + } + ) + invalid = pyarrow.table({"a": [1, 2], "b": [1, 2]}) + with pytest.raises(SchemaErrors) as exc_info: + schema.validate(invalid, lazy=True) + assert len(exc_info.value.failure_cases) >= 2 + + +def test_failure_cases_are_a_pyarrow_table(): + """failure_cases must come back as pyarrow, not depend on polars. + + The narwhals eager failure-case builder round-trips through polars. A + pyarrow install has no polars, so pyarrow gets its own builder — and the + reported type must not change based on whether polars happens to be + installed alongside. + """ + schema = pa.DataFrameSchema({"a": pa.Column(int, pa.Check.gt(0))}) + with pytest.raises(SchemaErrors) as exc_info: + schema.validate(pyarrow.table({"a": [1, -2, -3]}), lazy=True) + + failure_cases = exc_info.value.failure_cases + assert isinstance(failure_cases, pyarrow.Table) + assert set(failure_cases.column_names) >= { + "failure_case", + "schema_context", + "column", + "check", + "check_number", + "index", + } + assert failure_cases.column("failure_case").to_pylist() == ["-2", "-3"] + assert failure_cases.column("column").to_pylist() == ["a", "a"] + # index is null on the deferred-expr path, matching the polars backend + assert failure_cases.column("index").to_pylist() == [None, None] + + +def test_failure_cases_include_schema_level_errors(): + """Scalar (schema-level) and row-level failure cases must concatenate. + + They are built by different code paths; if they land on different + backends the concat step blows up. + """ + schema = pa.DataFrameSchema( + { + "a": pa.Column(int, pa.Check.gt(0)), + "missing": pa.Column(str), + } + ) + with pytest.raises(SchemaErrors) as exc_info: + schema.validate(pyarrow.table({"a": [1, -2]}), lazy=True) + + failure_cases = exc_info.value.failure_cases + assert isinstance(failure_cases, pyarrow.Table) + assert failure_cases.num_rows == 2 + assert set(failure_cases.column("failure_case").to_pylist()) == { + "missing", + "-2", + } + assert set(failure_cases.column("schema_context").to_pylist()) == { + "DataFrameSchema", + "Column", + } + + +def test_nullable_false_rejects_nulls(): + schema = pa.DataFrameSchema({"a": pa.Column(int, nullable=False)}) + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, None, 3]})) + + +def test_nullable_true_allows_nulls(): + schema = pa.DataFrameSchema({"a": pa.Column(int, nullable=True)}) + tbl = pyarrow.table({"a": [1, None, 3]}) + assert schema.validate(tbl).equals(tbl) + + +def test_unique_constraint(): + schema = pa.DataFrameSchema({"a": pa.Column(int, unique=True)}) + assert schema.validate(pyarrow.table({"a": [1, 2, 3]})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, 1, 3]})) + + +def test_column_coerce_is_not_supported(): + """Column-level coerce is a documented gap in the narwhals backend. + + pyarrow is served exclusively by the narwhals backends, where coercion is + deferred to v2 (see the strict xfail markers in + ``tests/narwhals/test_parity.py``). Assert the warned no-op behaviour so + this test flips loudly when coerce lands. + """ + schema = pa.DataFrameSchema({"a": pa.Column(str, coerce=True)}) + with pytest.warns(SchemaWarning, match="coerce=True is not applied"): + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [1, 2, 3]})) + + +def test_regex_column(): + schema = pa.DataFrameSchema({"^val_.+$": pa.Column(int, regex=True)}) + tbl = pyarrow.table({"val_a": [1], "val_b": [2], "other": ["x"]}) + assert schema.validate(tbl).equals(tbl) + + bad = pyarrow.table({"val_a": [1], "val_b": ["x"]}) + with pytest.raises(SchemaError): + schema.validate(bad) + + +def test_drop_invalid_rows(): + schema = pa.DataFrameSchema( + {"a": pa.Column(int, pa.Check.gt(0))}, drop_invalid_rows=True + ) + validated = schema.validate(pyarrow.table({"a": [1, -2, 3]}), lazy=True) + assert validated.column("a").to_pylist() == [1, 3] + + +def test_dataframe_level_check(): + schema = pa.DataFrameSchema( + {"a": pa.Column(int), "b": pa.Column(int)}, + checks=pa.Check( + lambda data: pc.less(data.table["a"], data.table["b"]), + ), + ) + assert schema.validate(pyarrow.table({"a": [1], "b": [2]})) is not None + with pytest.raises(SchemaError): + schema.validate(pyarrow.table({"a": [3], "b": [2]})) + + +def test_validation_disabled(table, schema): + invalid = pyarrow.table( + { + "int_col": [-1], + "float_col": [1.0], + "str_col": ["a"], + } + ) + with config_context(validation_enabled=False): + assert schema.validate(invalid) is invalid + + +def test_data_synthesis_not_supported(schema): + with pytest.raises(NotImplementedError): + schema.example() + with pytest.raises(NotImplementedError): + schema.strategy() diff --git a/tests/pyarrow/test_pyarrow_dtypes.py b/tests/pyarrow/test_pyarrow_dtypes.py new file mode 100644 index 000000000..0abe9a2ff --- /dev/null +++ b/tests/pyarrow/test_pyarrow_dtypes.py @@ -0,0 +1,83 @@ +"""Tests for dtype resolution in the pyarrow schema API.""" + +import pyarrow +import pytest + +import pandera.pyarrow as pa +from pandera.api.pyarrow.utils import pyarrow_dtype_to_narwhals, resolve_dtype +from pandera.errors import SchemaError + + +@pytest.mark.parametrize( + "dtype,expected", + [ + (pyarrow.int8(), "Int8"), + (pyarrow.int64(), "Int64"), + (pyarrow.uint16(), "UInt16"), + (pyarrow.float32(), "Float32"), + # pa.float64() stringifies as "double" — the string alias path + # cannot resolve it, so this covers the translation helper. + (pyarrow.float64(), "Float64"), + (pyarrow.bool_(), "Boolean"), + (pyarrow.string(), "String"), + (pyarrow.large_string(), "String"), + (pyarrow.date32(), "Date"), + ], +) +def test_pyarrow_dtype_translation(dtype, expected): + assert str(pyarrow_dtype_to_narwhals(dtype)) == expected + + +def test_parametrized_dtype_translation(): + assert str(pyarrow_dtype_to_narwhals(pyarrow.timestamp("us"))).startswith( + "Datetime" + ) + assert str( + pyarrow_dtype_to_narwhals(pyarrow.decimal128(10, 2)) + ).startswith("Decimal") + assert str(pyarrow_dtype_to_narwhals(pyarrow.list_(pyarrow.int32()))) == ( + "List(Int32)" + ) + + +@pytest.mark.parametrize("builtin", [int, float, str, bool]) +def test_python_builtins_resolve(builtin): + """Regression test: the narwhals engine used to reject python builtins.""" + assert resolve_dtype(builtin) is not None + + +@pytest.mark.parametrize( + "column_dtype,data,valid", + [ + (pyarrow.int64(), [1, 2, 3], True), + (pyarrow.float64(), [1.0, 2.0], True), + (pyarrow.float64(), [1, 2], False), + (int, [1, 2, 3], True), + (float, [1.0, 2.0], True), + (str, ["a"], True), + (bool, [True, False], True), + ("int64", [1, 2], True), + ("string", ["a"], True), + ], +) +def test_column_dtype_validation(column_dtype, data, valid): + schema = pa.DataFrameSchema({"a": pa.Column(column_dtype)}) + tbl = pyarrow.table({"a": data}) + if valid: + assert schema.validate(tbl).equals(tbl) + else: + with pytest.raises(SchemaError): + schema.validate(tbl) + + +def test_timestamp_column(): + import datetime + + schema = pa.DataFrameSchema({"ts": pa.Column(pyarrow.timestamp("us"))}) + tbl = pyarrow.table({"ts": pyarrow.array([datetime.datetime(2024, 1, 1)])}) + assert schema.validate(tbl).equals(tbl) + + +def test_unknown_dtype_raises(): + with pytest.raises(TypeError): + pa.Column(complex) diff --git a/tests/pyarrow/test_pyarrow_model.py b/tests/pyarrow/test_pyarrow_model.py new file mode 100644 index 000000000..283b884b9 --- /dev/null +++ b/tests/pyarrow/test_pyarrow_model.py @@ -0,0 +1,103 @@ +"""Tests for the pyarrow DataFrameModel API.""" + +import pyarrow +import pytest + +import pandera.pyarrow as pa +from pandera.errors import SchemaError +from pandera.typing.pyarrow import Table + + +class SimpleModel(pa.DataFrameModel): + int_col: int + str_col: str + + +def test_model_to_schema(): + schema = SimpleModel.to_schema() + assert isinstance(schema, pa.DataFrameSchema) + assert list(schema.columns) == ["int_col", "str_col"] + + +def test_model_validate(): + tbl = pyarrow.table({"int_col": [1, 2], "str_col": ["a", "b"]}) + assert SimpleModel.validate(tbl).equals(tbl) + + +def test_model_validate_wrong_dtype(): + tbl = pyarrow.table({"int_col": ["x"], "str_col": ["a"]}) + with pytest.raises(SchemaError): + SimpleModel.validate(tbl) + + +def test_model_with_field_checks(): + class Bounded(pa.DataFrameModel): + a: int = pa.Field(gt=0, le=10) + + assert Bounded.validate(pyarrow.table({"a": [1, 10]})) is not None + with pytest.raises(SchemaError): + Bounded.validate(pyarrow.table({"a": [0]})) + with pytest.raises(SchemaError): + Bounded.validate(pyarrow.table({"a": [11]})) + + +def test_model_optional_column(): + class WithOptional(pa.DataFrameModel): + a: int + b: str | None + + # 'b' is optional, so a table without it validates... + tbl = pyarrow.table({"a": [1]}) + assert WithOptional.validate(tbl).equals(tbl) + + # ...but a required column that is missing still fails. + class AllRequired(pa.DataFrameModel): + a: int + b: str + + with pytest.raises(SchemaError): + AllRequired.validate(tbl) + + +def test_model_nullable_field(): + class Nullable(pa.DataFrameModel): + a: int = pa.Field(nullable=True) + + tbl = pyarrow.table({"a": [1, None]}) + assert Nullable.validate(tbl).equals(tbl) + + +def test_model_custom_check(): + import pyarrow.compute as pc + + class WithCheck(pa.DataFrameModel): + a: int + + @pa.check("a") + @classmethod + def a_is_positive(cls, data): + return pc.greater(data.table[data.key], 0) + + assert WithCheck.validate(pyarrow.table({"a": [1, 2]})) is not None + with pytest.raises(SchemaError): + WithCheck.validate(pyarrow.table({"a": [-1]})) + + +def test_model_empty(): + empty = SimpleModel.empty() + assert empty.num_rows == 0 + assert empty.column_names == ["int_col", "str_col"] + assert empty.schema.field("int_col").type == pyarrow.int64() + + +def test_check_types_decorator(): + @pa.check_types + def transform(tbl: Table[SimpleModel]) -> Table[SimpleModel]: + return tbl + + valid = pyarrow.table({"int_col": [1], "str_col": ["a"]}) + assert transform(valid).equals(valid) + + invalid = pyarrow.table({"int_col": ["x"], "str_col": ["a"]}) + with pytest.raises(SchemaError): + transform(invalid) diff --git a/tests/pyarrow/test_pyarrow_register.py b/tests/pyarrow/test_pyarrow_register.py new file mode 100644 index 000000000..bc9f2cfed --- /dev/null +++ b/tests/pyarrow/test_pyarrow_register.py @@ -0,0 +1,69 @@ +"""Backend registration tests for the pyarrow API.""" + +import narwhals.stable.v1 as nw +import pyarrow +import pytest + +import pandera.pyarrow as pa +from pandera.api.checks import Check +from pandera.backends.narwhals.checks import NarwhalsCheckBackend +from pandera.backends.narwhals.components import ColumnBackend +from pandera.backends.narwhals.container import DataFrameSchemaBackend +from pandera.backends.pyarrow.register import register_pyarrow_backends + + +def test_pyarrow_table_resolves_narwhals_backends(): + register_pyarrow_backends() + tbl = pyarrow.table({"a": [1]}) + assert isinstance( + pa.DataFrameSchema.get_backend(tbl), DataFrameSchemaBackend + ) + assert isinstance(pa.Column.get_backend(tbl), ColumnBackend) + assert Check.get_backend(tbl) is NarwhalsCheckBackend + + +def test_importing_api_module_first_does_not_deadlock(): + """``pandera.api.pyarrow.*`` must be importable without the entry point. + + Registering eagerly from ``pandera.backends.pyarrow.__init__`` created a + circular import: the container imports the register module, whose package + ``__init__`` imported the container back. + """ + import subprocess + import sys + + for module in ( + "pandera.api.pyarrow.container", + "pandera.api.pyarrow.components", + "pandera.api.pyarrow.model", + ): + proc = subprocess.run( + [sys.executable, "-c", f"import {module}"], + capture_output=True, + text=True, + check=False, + ) + assert proc.returncode == 0, proc.stderr + + +def test_narwhals_check_dispatch_does_not_register_unrelated_backends(): + """A check on a narwhals frame must not pull in every backend. + + ``register_default_check_backends`` used to eagerly register the polars, + ibis and pyspark backends for any ``narwhals.*`` check object. That primed + each library's registration ``lru_cache``, so a later wipe of the shared + ``BACKEND_REGISTRY`` left them permanently unregistered. + """ + ibis_register = pytest.importorskip( + "pandera.backends.ibis.register" + ).register_ibis_backends + + register_pyarrow_backends() + # Ensure the narwhals check backends are present so the fast path applies. + Check.register_backend(nw.DataFrame, NarwhalsCheckBackend, force=True) + ibis_register.cache_clear() + + schema = pa.DataFrameSchema({"a": pa.Column(int, pa.Check.gt(0))}) + schema.validate(pyarrow.table({"a": [1, 2]})) + + assert ibis_register.cache_info().currsize == 0