Использование UDF Pandas Grouped Map с аргументами

Question

Я хочу использовать data.groupby.apply (), чтобы применить функцию к каждой строке моего кадра данных Pyspark в каждой группе.

Я использовал UDF Pandas сгруппированной карты. Однако я не могу понять, как добавить еще один аргумент в мою функцию.

Я пытался использовать аргумент в качестве глобальной переменной, но функция не распознает его (мой аргумент - фрейм данных pyspark)
Я также попробовал решения, предложенные в этом вопросе (для pandas dataframe) Используйте Pandas groupby () + apply () с аргументами

      @pandas_udf(schema,PandasUDFType.GROUPED_MAP)
        def function(key,data, interval):
            interval_df=interval.filter(interval["var"]==key).toPandas()
            for value in interval_df:
                  #Apply some operations

        return Data.groupBy("msn").apply(calc_diff, ('arg1'))

Или

 @pandas_udf(schema,PandasUDFType.GROUPED_MAP)
        def function(key,data, interval):
            interval_df=interval.filter(interval["var"]==key).toPandas()
            for value in interval_df:
                  #Apply some operations

        return Data.groupBy("msn").apply(lambda x: calc_diff(x,'arg1'))

Но я получаю ошибку:

ValueError: Неверная функция: pandas_udfs с типом функции GROUPED_MAP должен принимать либо один аргумент (данные), либо два аргумента (ключ, данные).

Может ли кто-нибудь помочь мне с вышеуказанным вопросом.

Спасибо

Answer 1

Я думаю, вы могли бы сделать что-то вроде этого

def myfun(data, key, interval):
    #Apply some operations
    return something

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def myfun_udf(data):
    return myfun(data=data, key=mykey, interval=myinterval)


mykey=1
myinterval=2

Data.groupBy("msn").apply(myfun_udf)