Writing checks for data input and output using Schema Model approach

[shoaibkhanz]

I am looking to validate a grp_by function which takes input and output schema. Trying to find a solution using a SchemaModel. Resources I have looked into is the following: resource#1: https://pandera.readthedocs.io/en/stable/schema_models.html#basic-usage and resource#2: https://pandera.readthedocs.io/en/stable/decorators.html?highlight=check_io#decorate-functions-and-coroutines. I am not sure If I have to write a 3rd class as mentioned in resource2.

Also is there a way to validate df_num_rows_check function when calling grp_by

I have intentionally misspelled the a column tranr_output which should be trans_output as per the schema, the checks fail to catch that, when I try something like the following.

@pa.check_types(lazy=True)
def grp_by(df: DataFrame[MyFirstSchema], trans_output: Series[Float]) -> Series[MyTransformedSchema]:
    df1 = df.groupby("VendorID")["trip_distance"].mean().reset_index()
    return df1.assign(tranr_output=df1["trip_distance"] * 100)

Any help would be appreciated..

from typing import Dict

import pandas as pd
import pandera as pa

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

@extensions.register_check_method()
def no_duplicates_in_df(df):
    return (df.duplicated().sum())==0


@extensions.register_check_method()
def records_before_2000(series):
    return (series < pd.Timestamp("2000-01-01")).sum() == 0


mylist = [1, 2, 5, 6]


class MyFirstSchema(pa.SchemaModel):
    VendorID: Series[int] = pa.Field(report_duplicates="all", isin=mylist)
    tpep_pickup_datetime: Series[pd.Timestamp] = pa.Field(
        lt=pd.Timestamp("2023-01-01"),
        nullable=False,
        in_range={
            "min_value": pd.Timestamp("2007-01-01"),
            "max_value": pd.Timestamp("2023-01-01"),
        },
        records_before_2000=(),
    )
    trip_distance: Series[float] = pa.Field(nullable=False)
    tip_amount: Series[float] = pa.Field(nullable=False)
    fare_amount: Series[float] = pa.Field(nullable=False)
    total_amount: Series[float] = pa.Field(nullable=False)

    @pa.dataframe_check
    @classmethod
    def tip_fare_bool(cls, df: pd.DataFrame) -> bool:
        return (df["tip_amount"] > df["fare_amount"]).sum() <= 20_000

    class Config:
        no_duplicates_in_df = ()
        strict = True

class MyTransformedSchema(MyFirstSchema):
    trans_output: Series[float] = pa.Field(nullable=False)

    @pa.dataframe_check
    @classmethod
    def df_num_rows_check(cls, series: pd.Series) -> bool:
        return series.shape[0] < 4

# how can I check input and output for the grp_by function outside of the MyTransformedSchema Class using SchemaModel Approach
@pa.check_io(df=MyFirstSchema, out=MyTransformedSchema)
def grp_by(df: DataFrame[MyFirstSchema]) -> Series[MyTransformedSchema]:
    df1 = df.groupby("VendorID")["trip_distance"].mean().reset_index()
    return df1.assign(tranr_output=df1["trip_distance"] * 100)


if __name__ == "__main__":
    # data link: https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2022-01.parquet
    df = pd.read_parquet("data/yellow_tripdata_2022-01.parquet")
    group_data = grp_by(df)