Union of schemas not working with pandera.pyspark

[benoitLebreton-perso]

Describe the bug
I started from this documentation example : https://pandera.readthedocs.io/en/stable/dataframe_models.html#validate-against-multiple-schemas which works in my env.

But when I'm adapting it with pandera.pyspark I have an error.

• I have checked that this issue has not already been reported.
• I have confirmed this bug exists on the latest version of pandera.
• (optional) I have confirmed this bug exists on the main branch of pandera.

Note: Please read this guide detailing how to provide the necessary information for us to reproduce your bug.

Code Sample, a copy-pastable example

from typing import Union
import pandera.pyspark as pa
from pandera.typing.pyspark_sql import DataFrame
from pyspark.sql import Row, SparkSession
from pyspark.sql.types import IntegerType


spark = SparkSession.builder.master("local").appName("test").getOrCreate()

class OnlyZeroesSchema(pa.DataFrameModel):
    a: IntegerType = pa.Field(eq=0)

class OnlyOnesSchema(pa.DataFrameModel):
    a: IntegerType = pa.Field(eq=1)

@pa.check_types
def return_zeros_or_ones(
    df: Union[DataFrame[OnlyZeroesSchema], DataFrame[OnlyOnesSchema]]
) -> Union[DataFrame[OnlyZeroesSchema], DataFrame[OnlyOnesSchema]]:
    return df


def create_df(spark: SparkSession, val: int):
    rows = [Row(a=val), Row(a=val)]
    df = spark.createDataFrame(rows)
    return df

df_0 = create_df(spark, 0)
# fails but it should pass
return_zeros_or_ones(df_0)
df_1 = create_df(spark, 1)
# fails but it should pass
return_zeros_or_ones(df_1)

The error is

Traceback (most recent call last):
  File "<string>", line 4, in <module>
  File "/Users/BENOIT/aifforecastsupply_ml_drift_demand_pyspark/.venv/lib/python3.9/site-packages/pandera/decorators.py", line 853, in _wrapper
    validated_pos, validated_kwd = validate_inputs(
  File "/Users/BENOIT/aifforecastsupply_ml_drift_demand_pyspark/.venv/lib/python3.9/site-packages/pandera/decorators.py", line 811, in validate_inputs
    validated_pos = validate_args(sig.bind_partial(*args).arguments, args)
  File "/Users/BENOIT/aifforecastsupply_ml_drift_demand_pyspark/.venv/lib/python3.9/site-packages/pandera/decorators.py", line 757, in validate_args
    return list(
  File "/Users/BENOIT/aifforecastsupply_ml_drift_demand_pyspark/.venv/lib/python3.9/site-packages/pandera/decorators.py", line 758, in <genexpr>
    _check_arg(arg_name, arg_value)
  File "/Users/BENOIT/aifforecastsupply_ml_drift_demand_pyspark/.venv/lib/python3.9/site-packages/pandera/decorators.py", line 711, in _check_arg
    raise errors.SchemaErrors(
pandera.errors.SchemaErrors: <exception str() failed>

Expected behavior

The expected behaviour should be same as in the doc : https://pandera.readthedocs.io/en/stable/dataframe_models.html#validate-against-multiple-schemas
The error should not happen

Desktop (please complete the following information):

• OS: iOS mac m1 macbook air
• Version: python 3.9 and pandera==0.20.4 with the extras pyspark (pandera[pyspark]==0.20.4)