[Python] How to return a complicated struct out of Pandas UDF in PySpark?

I struggle with writing a Pandas UDF that would return a complicated struct object. I have a function with the following signature:

def recipe_generator(
shop_type_column: pd.Series,
image_variant_id_column: pd.DataFrame,
product_id_column: pd.Series,
) -> pd.DataFrame:


where the return type of the struct is somewhat as follows:

return_type = T.StructType(
[
T.StructField("url", T.StringType()),
T.StructField(
"devicePixelRatios",
T.StructType(
[
T.StructField("original", T.StringType()),
T.StructField("1.5", T.StringType()),
T.StructField("2.0", T.StringType()),
]
),
),
T.StructField(
"formats",
T.StructType(
[
T.StructField("thumb-75", T.StringType()),
T.StructField("medium-200", T.StringType()),
T.StructField("medium-360", T.StringType()),
T.StructField("large-720", T.StringType()),
T.StructField("original", T.StringType()),
]
),
),
T.StructField("variants", T.ArrayType(T.StructType([T.StructField("id", T.StringType())]))),
]
)


However, I cannot achieve it - I am tinkering things here and there but to no avail. I was able to implement the solution via pure udf, the non-pandas version, but here - no success. For the record, that is the version of udf that worked:

F.udf(
lambda shop_type, image_variants, product_id: LinkCreator().prepare_image_url_recipe(
shop_type, image_variants, product_id
),
returnType=return_type,
)


Please bear in mind, that the image_variant_id_column is of array type, which, I believe, complicates the situation even more. Is there any error in the way I approach it? What should I correct to make it pass? Maybe there is some trick to enable return of such a complicated struct?

Continue reading...