Data Validation with Pandera

Flyte-SDK integrates with the Pandera library to provide robust data validation for DataFrames within your Flyte workflows. This integration allows you to define schemas for your pandas, Polars, and PySpark SQL DataFrames, ensuring data quality and consistency as data flows through your tasks. When validation fails, Flyte automatically generates rich HTML reports that provide detailed insights into the errors, making debugging and data quality monitoring straightforward.

Enabling Data Validation

The flyteplugins.pandera plugin automatically detects and registers the necessary type transformers for pandas, Polars, and PySpark SQL DataFrames if their respective libraries are installed in your environment. This conditional registration is handled by the flyteplugins.pandera.register.register_type_transformers() function.

For example, if pandas is installed, the flyteplugins.pandera.transformers.pandas.PanderaPandasDataFrameTransformer is registered with Flyte's TypeEngine, enabling automatic validation for pandas DataFrames. Similar transformers exist for Polars (flyteplugins.pandera.transformers.polars.PanderaPolarsDataFrameTransformer) and PySpark SQL (flyteplugins.pandera.transformers.pyspark_sql.PanderaPySparkSqlDataFrameTransformer).

Defining Data Schemas

You define your data expectations using Pandera's schema definition capabilities. This typically involves creating a class that inherits from pandera.DataFrameModel and specifying the columns and their data types using pandera.typing.Series.

Here's an example of a schema for a pandas.DataFrame:

import pandera as pa
from pandera.typing import DataFrame, Series

class MyPandasSchema(pa.DataFrameModel):
    feature_1: Series[int] = pa.Field(ge=0)
    feature_2: Series[float] = pa.Field(gt=0, le=1)
    category: Series[str]

For Polars DataFrames, you would use pandera.typing.polars.DataFrame and pandera.typing.polars.Series:

import pandera as pa
import polars as pl
from pandera.typing.polars import DataFrame, Series

class MyPolarsSchema(pa.DataFrameModel):
    id: Series[int]
    value: Series[float]
    timestamp: Series[pl.Datetime]

And for PySpark SQL DataFrames, you use pandera.typing.pyspark_sql.DataFrame and pandera.typing.pyspark_sql.Series:

import pandera as pa
from pandera.typing.pyspark_sql import DataFrame, Series

class MyPySparkSchema(pa.DataFrameModel):
    user_id: Series[int]
    event_name: Series[str]
    event_time: Series[pa.Timestamp]

Applying Validation in Tasks

Once you have defined your Pandera schema, you can apply it directly to the input and output type annotations of your Flyte tasks. Flyte-SDK's Pandera plugin automatically intercepts the DataFrame objects, validates them against the specified schema, and handles any validation failures. The core logic for this is found in the _validate method of the flyteplugins.pandera.transformers.base.PanderaDataFrameTransformer class and its subclasses.

Here's how you would use the MyPandasSchema in a Flyte task:

import pandas as pd
import pandera as pa
from pandera.typing import DataFrame
from flytekit import task, workflow

class MyPandasSchema(pa.DataFrameModel):
    feature_1: Series[int] = pa.Field(ge=0)
    feature_2: Series[float] = pa.Field(gt=0, le=1)
    category: Series[str]

@task
def process_data(df: DataFrame[MyPandasSchema]) -> DataFrame[MyPandasSchema]:
    # The input 'df' is automatically validated against MyPandasSchema
    # Any DataFrame returned will also be validated against MyPandasSchema
    df['new_feature'] = df['feature_1'] * df['feature_2']
    return df

@workflow
def my_workflow(input_df: DataFrame[MyPandasSchema]) -> DataFrame[MyPandasSchema]:
    return process_data(df=input_df)

# Example of how to run locally (not part of Flyte task definition)
if __name__ == "__main__":
    valid_data = pd.DataFrame({
        "feature_1": [1, 2, 3],
        "feature_2": [0.1, 0.5, 0.9],
        "category": ["A", "B", "C"]
    })
    # This would pass validation
    # my_workflow(input_df=valid_data)

    invalid_data = pd.DataFrame({
        "feature_1": [-1, 2, 3], # Invalid: ge=0
        "feature_2": [0.1, 1.5, 0.9], # Invalid: le=1
        "category": ["A", "B", "C"]
    })
    # This would raise a SchemaError by default
    try:
        my_workflow(input_df=invalid_data)
    except pa.errors.SchemaErrors as e:
        print(f"Validation failed as expected: {e}")

Customizing Validation Behavior

By default, if a DataFrame fails validation against its Pandera schema, Flyte-SDK will raise a pandera.errors.SchemaErrors exception, stopping the task execution. You can customize this behavior using the flyteplugins.pandera.ValidationConfig class and typing_extensions.Annotated.

The ValidationConfig dataclass has a single attribute, on_error, which can be set to:

"raise" (default): Raises a SchemaErrors exception on validation failure.
"warn": Logs a warning message and allows the task to continue, returning the original (unvalidated) DataFrame.

To apply a custom ValidationConfig, you wrap your schema annotation with Annotated:

from typing_extensions import Annotated
import pandas as pd
import pandera as pa
from pandera.typing import DataFrame, Series
from flytekit import task, workflow
from flyteplugins.pandera import ValidationConfig

class MyLenientSchema(pa.DataFrameModel):
    col_a: Series[int]
    col_b: Series[str]

@task
def lenient_validation_task(
    df: Annotated[DataFrame[MyLenientSchema], ValidationConfig(on_error="warn")]
) -> DataFrame[MyLenientSchema]:
    # Validation errors will be logged as warnings, and the task will proceed.
    print("Task received DataFrame, validation warnings (if any) were logged.")
    return df

@workflow
def lenient_workflow(input_df: DataFrame[MyLenientSchema]) -> DataFrame[MyLenientSchema]:
    return lenient_validation_task(df=input_df)

# Example of how to run locally
if __name__ == "__main__":
    invalid_data_lenient = pd.DataFrame({
        "col_a": [1, "not_an_int", 3], # Invalid type
        "col_b": ["x", "y", "z"]
    })
    print("Running lenient workflow with invalid data...")
    lenient_workflow(input_df=invalid_data_lenient)
    print("Lenient workflow completed, check logs for warnings.")

Internally, the flyteplugins.pandera.transformers.base._extract_config helper function is responsible for extracting the ValidationConfig from the type annotations, and the _validate method then applies the specified on_error behavior.

Understanding Validation Reports

Whether validation succeeds or fails, Flyte-SDK's Pandera plugin automatically generates detailed HTML reports. These reports are accessible in the Flyte UI and provide a comprehensive overview of the validation process.

For successful validations, the report includes a summary and a preview of the validated data. For failures, the report highlights:

Summary: High-level metadata about the schema and total errors.
Data Preview: A sample of the data that was validated.
Schema-level Errors: Issues related to the schema definition itself, such as missing columns or incorrect data types.
Data-level Errors: Problems with the data values, such as nulls, out-of-range values, or values not matching specified patterns.

These reports are generated by implementations of the flyteplugins.pandera.renderers.base.PanderaReportRenderer protocol. Specifically:

flyteplugins.pandera.renderers.pandas.PanderaPandasReportRenderer for pandas DataFrames.
flyteplugins.pandera.renderers.polars.PanderaPolarsReportRenderer for Polars DataFrames.
flyteplugins.pandera.renderers.pyspark_sql.PanderaPySparkSqlReportRenderer for PySpark SQL DataFrames.

These renderers leverage the great_tables library to produce visually rich and informative HTML output. For Polars and PySpark SQL DataFrames, the respective renderers (PanderaPolarsReportRenderer and PanderaPySparkSqlReportRenderer) convert a limited subset of the DataFrame to a pandas.DataFrame for rendering purposes, ensuring compatibility with great_tables while managing potential memory constraints for large datasets.

Supported DataFrame Types

The flyteplugins.pandera plugin provides native support for data validation across the following DataFrame libraries:

pandas: Use pandera.typing.DataFrame and pandera.typing.Series with pandas.DataFrame objects.
Polars: Use pandera.typing.polars.DataFrame or pandera.typing.polars.LazyFrame with polars.DataFrame or polars.LazyFrame objects.
PySpark SQL: Use pandera.typing.pyspark_sql.DataFrame with pyspark.sql.DataFrame objects.

This broad support ensures that you can maintain data quality standards regardless of your chosen DataFrame processing engine within Flyte. The flyteplugins.pandera.transformers.base.PanderaDataFrameTransformer serves as the abstract base for these specific implementations, providing a consistent validation mechanism across different DataFrame types. Its _resolve_native_df_type method is implemented by each subclass to correctly identify the underlying DataFrame type being validated.

Enabling Data Validation​

Defining Data Schemas​

Applying Validation in Tasks​

Customizing Validation Behavior​

Understanding Validation Reports​