Передача столбца фрейма данных и внешнего списка в udf под withColumn

У меня есть датафрейм Spark со следующей структурой. BodyText_token имеет токены (обработано/набор слов). И у меня есть вложенный список определенных ключевых слов

root
 |-- id: string (nullable = true)
 |-- body: string (nullable = true)
 |-- bodyText_token: array (nullable = true)

keyword_list=['union','workers','strike','pay','rally','free','immigration',],
['farmer','plants','fruits','workers'],['outside','field','party','clothes','fashions']]

Мне нужно было проверить, сколько токенов попадает под каждый список ключевых слов, и добавить результат в виде нового столбца существующего фрейма данных. Например: если tokens =["become", "farmer","rally","workers","student"] результат будет → [1,2,0]

Следующая функция работала как ожидалось.

def label_maker_topic(tokens,topic_words):
    twt_list = []
    for i in range(0, len(topic_words)):
        count = 0
        #print(topic_words[i])
        for tkn in tokens:
            if tkn in topic_words[i]:
                count += 1
        twt_list.append(count)

    return twt_list

Я использовал udf в withColumn для доступа к функции и получаю сообщение об ошибке. Я думаю о передаче внешнего списка в udf. Можно ли как-то передать внешний список и столбец datafram в udf и добавить новый столбец в мой dataframe?

topicWord = udf(label_maker_topic,StringType())
myDF=myDF.withColumn("topic_word_count",topicWord(myDF.bodyText_token,keyword_list))

Ответ 1

Самое чистое решение - передать дополнительные аргументы, используя закрытие:

def make_topic_word(topic_words):
     return udf(lambda c: label_maker_topic(c, topic_words))

df = sc.parallelize([(["union"], )]).toDF(["tokens"])

(df.withColumn("topics", make_topic_word(keyword_list)(col("tokens")))
    .show())

Это не требует изменений в keyword_list или функции, которую вы обертываете с помощью UDF. Вы также можете использовать этот метод для передачи произвольного объекта. Это можно использовать для передачи, например, списка sets для эффективного поиска.

Если вы хотите использовать свой текущий UDF и передать topic_words напрямую, вам нужно сначала преобразовать его в литерал столбца:

from pyspark.sql.functions import array, lit

ks_lit = array(*[array(*[lit(k) for k in ks]) for ks in keyword_list])
df.withColumn("ad", topicWord(col("tokens"), ks_lit)).show()

В зависимости от ваших данных и требований могут быть альтернативные, более эффективные решения, которые не требуют UDF (разбить + агрегат + сбой) или поисковые запросы (хеширование + векторные операции).

Ответ 2

Следующее работает отлично, когда любой внешний параметр может быть передан в UDF (измененный код, чтобы помочь кому-либо)

topicWord=udf(lambda tkn: label_maker_topic(tkn,topic_words),StringType())
myDF=myDF.withColumn("topic_word_count",topicWord(myDF.bodyText_token))

Ответ 3

Другой способ - использовать частичное из модуля functools

from functools import partial

func_to_call = partial(label_maker_topic, topic_words=keyword_list)

pyspark_udf = udf(func_to_call, <specify_the_type_returned_by_function_here>)

df = sc.parallelize([(["union"], )]).toDF(["tokens"])

df.withColumn("topics", pyspark_udf(col("tokens"))).show()