Функция передачи Pyspark как параметр в UDF - PullRequest
Новая
       106

Функция передачи Pyspark как параметр в UDF

0 голосов
/ 05 августа 2020

Я пытаюсь создать UDF, который принимает в качестве параметра другую функцию. Но казнь заканчивается исключением. Код, который я запускаю: 

import pandas as pd
from pyspark import SparkConf, SparkContext, SQLContext
from pyspark.sql.types import MapType, DataType, StringType
from pyspark.sql.functions import udf, struct, lit
import os

sc = SparkContext.getOrCreate(conf=conf)
sqlContext = SQLContext(sc)

df_to_test = sqlContext.createDataFrame(
    pd.DataFrame({
        'inn': ['111', '222', '333'],
        'field1': [1, 2, 3],
        'field2': ['a', 'b', 'c']
    }))

def foo_fun(row, b) -> str:
    return 'a' + b()

def bar_fun():
    return 'I am bar'

foo_fun_udf = udf(foo_fun, StringType())
df_to_test.withColumn(
    'foo', 
    foo_fun_udf(struct([df_to_test[x] for x in df_to_test.columns]), bar_fun)
).show()

Исключение: 

Invalid argument, not a string or column: <function bar_fun at 0x7f0e69ce6268> of type <class 'function'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function.

Я безуспешно пытался обернуть bar_fun в udf. Есть ли способ передать функцию в качестве параметра?

1 Ответ

1 голос
/ 05 августа 2020

Вы не так уж далеки от решения. Вот как бы я это сделал: 

def foo_fun_udf(func):

    def foo_fun(row) -> str:
        return 'a' + func()

    out_udf = udf(foo_fun, StringType())
    return out_udf 

df_to_test.withColumn(
    'foo', 
    foo_fun_udf(bar_fun)(struct([df_to_test[x] for x in df_to_test.columns]))
).show()
...