Pandera: Towards Better Data Testing Tools for Data Science and Machine Learning

Niels Bantilan

Scipy, July 16th 2021

This talk about the road to better data testing tools for data science and machine learning

Outline 📝

• Introduction to Data Testing

• Pandera Quickstart

• Roadmap: Guiding Principles

• How to Contribute

Introduction to Data Testing

"Data Testing" is the colloquial term [for] "schema validation" or "data validation"... It's merely a fancy way of [saying], "are my data as I expect them to be?" - Eric Ma

Addendum: "Data Testing" can also be thought of as testing the transformation code that produces the data.

To give you a simple example...

A Simple Example: Life Before Pandera

data_cleaner.py

import pandas as pd

raw_data = pd.DataFrame({
    "continuous": ["-1.1", "4.0", "10.25", "-0.1", "5.2"],
    "categorical": ["A", "B", "C", "Z", "X"],
})

def clean(raw_data):
    return (
        raw_data
        # do some cleaning 🧹✨
    )

test_data_cleaner.py

def test_clean():
    mock_raw_data = ...  # hand-written mock data 😅
    result = clean(mock_raw_data)

    # assumptions about clean data
    assert result["continuous"].ge(0).all()
    assert result["categorical"].isin(["A", "B", "C"]).all()

Pandera Quickstart

An expressive and light-weight statistical validation tool for dataframes

• Check the types and properties of dataframes

• Easily integrate with existing data pipelines via function decorators

• Synthesize data from schema objects for property-based testing

Object-based API

Defining a schema looks and feels like defining a pandas dataframe

import pandera as pa

schema = pa.DataFrameSchema(
    columns={
        "continuous": pa.Column(float, pa.Check.ge(0)),
        "categorical": pa.Column(str, pa.Check.isin(["A", "B", "C"])),
    },
    coerce=True,
)

Class-based API

Inspired by pydantic

from pandera.typing import Series

class Schema(pa.SchemaModel):
    continuous: Series[float] = pa.Field(ge=0)
    categorical: Series[str] = pa.Field(isin=["A", "B", "C"])

    class Config:
        coerce = True

Pandera comes in two flavors

Pandera Parses and Validates Data

With coerce=True pandera first performs type coercion on the columns before validating them.

raw_data = pd.DataFrame({
    "continuous": ["-1.1", "4.0", "10.25", "-0.1", "5.2"],
    "categorical": ["A", "B", "C", "Z", "X"],
})

try:
    Schema.validate(raw_data, lazy=True)
except pa.errors.SchemaErrors as exc:
    display(exc.failure_cases)

	schema_context	column	check	check_number	failure_case	index
0	Column	continuous	greater_than_or_equal_to(0)	0	-1.1	0
1	Column	continuous	greater_than_or_equal_to(0)	0	-0.1	3
2	Column	categorical	isin({'B', 'A', 'C'})	0	Z	3
3	Column	categorical	isin({'B', 'A', 'C'})	0	X	4

The core API of pandera

As a meta-point, this presentation is built with jupyter, so almost all the code in this presentation is validated and tested with pandera

🛣 Roadmap: Guiding Principles

I'll outline pandera's roadmap by describing to you four principles that guide the project's development. Along the way I'll highlight a few of pandera's features and specific roadmap items that are it the repo's issues page.

Principle 1: Parse, then Validate

pydantic [and pandera guarantee] the types and constraints of the output [data], not the input data. -Pydantic Docs

raw_data = ...
valid_data = validate(parse(raw_data))  # raise Exception if constraints are not met

Pandera guarantees that input and output dataframes fulfill the types and constraints as defined by type annotations.

raw_data = pd.DataFrame({
    "continuous": list("123456"),
    "categorical": list("AABBCC"),
})

class Schema(pa.SchemaModel):
    continuous: Series[float] = pa.Field(ge=0)
    categorical: Series[str] = pa.Field(isin=["A", "B", "C"])

    class Config:
        coerce = True

from pandera.typing import DataFrame

@pa.check_types
def summarize_data(clean_data: DataFrame[Schema]):
    return clean_data.groupby("categorical")["continuous"].mean()

display(summarize_data(raw_data).rename("mean_continuous").to_frame())

	mean_continuous
categorical
A	1.5
B	3.5
C	5.5

🛣 Roadmap Item

Extend parsing functionality to support arbitrary transformations

class Schema(pa.SchemaModel):
    continuous: Series[float] = pa.Field(ge=0)
    categorical: Series[str] = pa.Field(isin=["A", "B", "C"])

    class Config:
        coerce = True

    @pa.parser("continuous")
    def truncate_continuous(cls, series):
        """set negative values to nan"""
        return series.mask(series < 0, pd.NA)

    @pa.parser("continuous")
    def filter_continuous(cls, series):
        """filter out records with negative values in the continuous column"""
        return series[series >= 0]

Principle 2: Make Schemas Reuseable, Adaptable, and Portable

Once you've defined a schema, you can use it in your source code

# data_cleaning.py
from pandera.typing import DataFrame

@pa.check_types
def clean_data(raw_data) -> DataFrame[Schema]:
    return (
        raw_data
        # do some cleaning
    )

... and your test suite (or anywhere you want, really!)

# test_data_cleaning.py
def test_clean_data():
    raw_data = ...
    clean_data(raw_data)

Now the output dataframe type is validated when you call clean_data at runtime so our test reduces to an execution test!

Adaptability: define a base schema and build on top of it

class InputSchema(pa.SchemaModel):
    _categories = ["A", "B", "C"]  # store arbitrary metadata in private class attributes
    continuous: Series[float] = pa.Field(ge=0)
    categorical: Series[str] = pa.Field(isin=_categories)

    class Config:
        coerce = True

class OutputSchema(InputSchema):
    categorical_one_hot: Series[int] = pa.Field(alias="one_hot_", regex=True)

    @pa.check("one_hot_")
    def categorical_one_hot_check(cls, series):
        return series.name[-1] in cls._categories

@pa.check_types
def featurize_data(clean_data: DataFrame[InputSchema]) -> DataFrame[OutputSchema]:
    one_hot = pd.get_dummies(clean_data["categorical"], prefix="one_hot")
    return pd.concat([clean_data, one_hot], axis="columns")

display(featurize_data(raw_data).head(3))

	continuous	categorical	one_hot_A	one_hot_B	one_hot_C
0	1.0	A	1	0	0
1	2.0	A	1	0	0
2	3.0	B	0	1	0

Since, dataframes are complex objects, pandera focuses on making the process of defining schemas as concise as possible, offloading the concerns around column types and allowable values so you can focus more on the analysis/modeling logic.

Portability: Support Other Dataframe Libraries and Schema Specifications in the Ecosystem

Support frictionless data table schemas (✨ coming out in the 0.7.0 release ✨)

from pandera.io import from_frictionless_schema

frictionless_schema = {
    "fields": [
        {
            "name": "column_1",
            "type": "integer",
            "constraints": {"minimum": 10, "maximum": 99}
        }
    ],
    "primaryKey": "column_1"
}

pandera_schema = from_frictionless_schema(frictionless_schema)

🛣 Roadmap Items

• ✅ Decouple pandera and pandas type systems

• Abstract out parsing/validation logic to support non-pandas dataframes

• Add Titles and Descriptions for SchemaModels

Principle 3: Generative Schemas Facilitate Property-based Testing

You have a schema with a bunch of metadata about it... why not generate data for testing?

display(InputSchema.example(size=3))

	continuous	categorical
0	1.100000e+00	B
1	1.192093e-07	B
2	9.007199e+15	B

input_schema_strategy = InputSchema.strategy(size=5)
print(input_schema_strategy)
print(type(input_schema_strategy))

_dataframe_strategy()
<class 'hypothesis.strategies._internal.lazy.LazyStrategy'>

from hypothesis import given

@given(input_schema_strategy)
def test_featurize_data(clean_data):
    featurize_data(clean_data)

test_featurize_data()

Automate the tedium of hand-writing mock dataframes for testing!

Generate schemas as multi-purpose artifacts

At run-time

graph LR R[(raw data)] -- raw schema --> P[data processor] P -- clean schema --> C[(clean data)] C -- clean schema --> F[featurizer] F -- feature schema --> X[(feature data)]

At test-time

graph LR S(raw schema) --> R[(mock raw data)] R --> P[data processor] P -- clean schema --> O[output]

graph LR S(clean schema) --> R[(mock clean data)] R --> F[featurizer] F -- feature schema --> O[output]

🛣 Roadmap Items

• Add global schema-level override strategy

• Support data synthesis strategies for joint distributions

• Make Hypothesis strategies more efficient

Principle 4: Profile Data and Data Pipelines

Pandera uses basic data profiling to infer a schema from realistic data

realistic_data = pd.DataFrame({"continuous": [1, 2, 3, 4, 5, 6]})
bootstrapped_schema = pa.infer_schema(realistic_data)
print(bootstrapped_schema)

<Schema DataFrameSchema(
    columns={
        'continuous': <Schema Column(name=continuous, type=int64)>
    },
    checks=[],
    coerce=True,
    pandas_dtype=None,
    index=<Schema Index(name=None, type=int64)>,
    strict=False
    name=None,
    ordered=False
)>

Write it out into a python file with bootstrapped_schema.to_script("schemas.py")

from pandera import (
    DataFrameSchema,
    Column,
    Check,
    Index,
    MultiIndex,
    PandasDtype,
)

schema = DataFrameSchema(
    columns={
        "continuous": Column(
            pandas_dtype=PandasDtype.Int64,
            checks=[
                Check.greater_than_or_equal_to(min_value=1.0),
                Check.less_than_or_equal_to(max_value=6.0),
            ],
            nullable=False,
            allow_duplicates=True,
            coerce=False,
            required=True,
            regex=False,
        )
    },
    index=Index(
        pandas_dtype=PandasDtype.Int64,
        checks=[
            Check.greater_than_or_equal_to(min_value=0.0),
            Check.less_than_or_equal_to(max_value=5.0),
        ],
        nullable=False,
        coerce=False,
        name=None,
    ),
    coerce=True,
    strict=False,
    name=None,
)

Write it out into a yaml file with bootstrapped_schema.to_yaml("schema.yaml")

schema_type: dataframe
version: 0.6.5
columns:
  continuous:
    pandas_dtype: int64
    nullable: false
    checks:
      greater_than_or_equal_to: 1.0
      less_than_or_equal_to: 6.0
    allow_duplicates: true
    coerce: false
    required: true
    regex: false
checks: null
index:
- pandas_dtype: int64
  nullable: false
  checks:
    greater_than_or_equal_to: 0.0
    less_than_or_equal_to: 5.0
  name: null
  coerce: false
coerce: true
strict: false

🛣 Roadmap Item

Create schema from a pandas-profiling ProfileReport

from pandera.io import from_pandas_profile_report

df = ...
profile = ProfileReport(df)
schema = from_pandas_profile_report(profile)

Profile Data and Data Pipelines

If we have pandera schema type annotations...

import pandera as pa
from pandera.typing import DataFrame as DF

from schemas import Raw, Clean, Training

@pa.check_types
def load() -> DF[Raw]:
    ...

@pa.check_types
def clean(raw_data: DF[Raw]) -> DF[Clean]:
    ...

@pa.check_types
def featurize(clean_data: DF[Clean]) -> DF[Training]:
    ...

@pa.check_types
def train_model(training_data: DF[Training]):
    ...

We can potentially create a data flow graph

graph LR L[load] --DF-Raw--> C[clean] C --DF-Clean--> F[featurize] F --DF-Training--> T[train_model]

Collect coverage statistics of schema-annotated dataframes to identify points in the pipeline that lack dataframe type coverage

Function	Input Type	Output Type	Test Errors
load	-	DF[Raw]	0
clean	DF[Raw]	DF[Clean]	1
featurize	DF[Clean]	DF[Train]	7
train_model	DF[Train]	-	2

🛣 Roadmap Items

• Parse schema-decorated functions to construct a dataflow graph

• Implement error report aggregator

• Implement CLI for pipeline profiling and reports

• Create pytest-pandera plugin for profiling data pipelines in your test suite

🛣 Roadmap

Guiding Principle	Description	Issue
Parse, then Validate	Extend parsing functionality to support arbitrary transformations
Make Schemas Reusable, Adaptable, and Portable	Support Other Schema Specifications in the Ecosystem
	Decouple pandera and pandas type systems
	Abstract out parsing/validation logic to support non-pandas dataframes
	Add Titles and Descriptions for SchemaModels
Generative Schemas Facilitate Property-based Testing	Add global schema-level override strategy
	Support data synthesis strategies for joint distributions
	Make Hypothesis strategies more efficient
Profile Data and Data Pipelines	Create schema from a pandas-profiling ProfileReport
	Parse schema-decorated functions to construct a dataflow graph
	Implement error report aggregator
	Implement CLI for pipeline profiling and reports
	Create pytest-pandera plugin for profiling data pipelines in your test suite

Where to Learn More

• Pycon [2021] - Statistical Typing: A Runtime TypingSystem for Data Science and Machine Learning
  • video: https://youtu.be/PI5UmKi14cM
• Scipy [2020] - Statistical Data Validation of Pandas Dataframes
  • video: https://youtu.be/PxTLD-ueNd4
  • talk: https://conference.scipy.org/proceedings/scipy2020/pdfs/niels_bantilan.pdf
• Pandera Blog [2020]: https://blog.pandera.ci/statistical%20typing/unit%20testing/2020/12/26/statistical-typing.html
• PyOpenSci Blog [2019]: https://www.pyopensci.org/blog/pandera-python-pandas-dataframe-validation
• Personal Blog [2018]: https://cosmicbboy.github.io/2018/12/28/validating-pandas-dataframes.html

How to Contribute

• Repo: https://github.com/pandera-dev/pandera
• Docs: https://pandera.readthedocs.io
• Contributing Guide: https://pandera.readthedocs.io/en/stable/CONTRIBUTING.html

Join the Scipy Mentored Sprints! 👟👟

Toss a coin to your maintainer 👍🪙 https://github.com/sponsors/cosmicBboy

🎉 Shoutouts to pyopensci all the pandera contributors! 🎉