How to pass variable to UDAF (user define aggregation function) - pyspark

Question

import pandas as pd
import numpy as np
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from pyspark.sql.functions import PandasUDFType, pandas_udf
from pyspark.sql.types import *
import os
#pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
def split(df, validation_period):
""Logic""
return df
def train_test_split(spark, data_frame, request_json_data):
data_frame = spark.createDataFrame(data_frame)
print(data_frame.schema)
validation_period = request_json_data['validation_period']
groupby_key = request_json_data['groupby_key']
data_frame.groupby(groupby_key).apply(split, validation_period).show()
Can't call split function it gives error. apply() takes 2 positional arguments but 3 were given. I want to pass validation_period as argument to split function.

Answer 1

Short answer: you can't pass extra argument to pandas grouped map udf since it excepts only a single pandas df as an argument.
Long answer: There are other ways you can pass validation_period to the function
use some form of a closure
def split_fabric(validation_period):
#pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
def split(df):
""Logic""
return df
pass it as a column
data_frame \
.withColumn("validation_period", F.lit(validation_period)) \
.groupby(groupby_key).apply(split, validation_period).show()