Support for multiple argument dispatch functions based on schema validation

[billyvinning]

Hi,

Firstly, I think that pandera is a great tool for writing production-ready pandas projects. Where I think there is room for extension in pandera is in its data pipelining integration. One nice feature to have when defining data pipelines, I think, would be multiple dispatch functions.

There are a few ways of achieving multiple dispatch functions with traditional Python types, these being functools.singledispatch, multipledispatch and multimethod. I am proposing an analogous support in pandera, where the dispatch function is matched based on validation of the data passed. The ultimate use case for this is for the user to be able to write generic data pipelines that accept multiple schemas and validate the inputs and outputs of each step. As a minimal example, put explicitly, given that we have some schema defined: A, B and C, we want to be able to define a single interface that a) transforms both A->C and B->C and b) validates the inputs and outputs.

Has this idea ever been considered for pandera? If not, is there any interest for this kind of thing?

I was able to achieve a dirty minimal prototype by extending multimethod.multimethod, please see below for the code. The basic premise of this snippet is that my implementation of multischema.register converts occurrences of pa.typing.DataFrame[...] in the type hints of registered functions into multimethod.parametric. My usage of multimethod.parametric defines a custom type tied to each schema that defines isinstance based on whether DataFrameModel.validate(df) throws an exception or not. I recognise that we would have to consider what should happen when some given data can be validated against multiple schema that are registered by the dispatch function.

Environment:

python=3.11.7
multimethod=1.11.2
pandas=2.2.1
pandera=0.18.3

Example:

import functools
import types
from typing import Any, Callable, get_args, get_origin, overload

import pandera as pa
from multimethod import REGISTERED, multimethod, parametric
from pandera.typing import DataFrame, Series


class Schema(pa.DataFrameModel):
    age_in_weeks: Series[int]


class AnotherSchema(pa.DataFrameModel):
    age_in_days: Series[int]


class SchemaOut(pa.DataFrameModel):
    age_in_years: Series[int]

    class Config:
        strict = True


def pa_generic_as_parametric(generic: type) -> Callable[[DataFrame[Any]], bool]:
    """Get DataFrame[...] as a custom type that runs `DataFrameModel.validate(...)`
    as the `isinstance` check.
    """
    schema = get_args(generic)
    if len(schema) > 1:
        # Assume generic has only one schema in this minimal example.
        raise NotImplementedError

    def follows_schema(df: DataFrame[Any]) -> bool:
        try:
            schema[0].validate(df)
        except pa.errors.SchemaError:
            return False
        return True

    return parametric(DataFrame, follows_schema)


def regulate_pandera_hinted_function(f: Callable[..., Any]) -> Callable[..., Any]:
    """If function contains a DataFrame[...] generic then copy the function
    and convert the typehint to `multimethod.parametric`, else return original function.
    """
    if not any(get_origin(v) is DataFrame for v in f.__annotations__.values()):
        return f
    g = types.FunctionType(
        f.__code__,
        f.__globals__,
        name=f.__name__,
        argdefs=f.__defaults__,
        closure=f.__closure__,
    )
    g = functools.update_wrapper(g, f)
    g.__kwdefaults__ = f.__kwdefaults__
    g.__annotations__ = {
        k: pa_generic_as_parametric(v) if get_origin(v) is DataFrame else v
        for k, v in f.__annotations__.items()
    }
    return g


class multischema(multimethod):
    """Extend `multimethod.multimethod` single dispatch style to work with `pandera` type hints."""

    @overload
    def register(self, __func: REGISTERED) -> REGISTERED: ...  # pragma: no cover

    @overload
    def register(
        self,
        *args: type,
    ) -> Callable[[REGISTERED], REGISTERED]: ...  # pragma: no cover

    def register(self, *args: Any) -> Callable[..., Any]:
        args = tuple([regulate_pandera_hinted_function(arg) for arg in args])
        return super().register(*args)


@multischema
def transform(df: DataFrame[Any]):
    """Single dispatch function."""


@transform.register
def _(df: DataFrame[Schema]) -> DataFrame[SchemaOut]:
    """Implementation of single dispatch function for Schema."""
    age_in_years = df[Schema.age_in_weeks] // 52
    return age_in_years.to_frame(name=SchemaOut.age_in_years).pipe(
        DataFrame[SchemaOut],
    )


@transform.register
def _(df: DataFrame[AnotherSchema]) -> DataFrame[SchemaOut]:
    """Implementation of single dispatch function for AnotherSchema."""
    age_in_years = df[AnotherSchema.age_in_days] // 365
    return age_in_years.to_frame(name=SchemaOut.age_in_years).pipe(
        DataFrame[SchemaOut],
    )


def main() -> None:
    df_1 = DataFrame[Schema]({Schema.age_in_weeks: [1865, 1500, 2900]})
    df_2 = DataFrame[AnotherSchema]({AnotherSchema.age_in_days: [4000, 8053, 5555]})

    df_1_out = transform(df_1)
    df_2_out = transform(df_2)

    print(df_1_out)
    print(df_2_out)


if __name__ == "__main__":
    main()

Expected output:

python script.py
  age_in_years
0            35
1            28
2            55
   age_in_years
0            10
1            22
2            15
mypy script.py
Success: no issues found in 1 source file

Let me know what you think!