web-dev-qa-db-fra.com

Passer une colonne de trame de données et une liste externe à udf sous avecColonne

J'ai un Spark dataframe avec la structure suivante. Le bodyText_token a les jetons (traités/ensemble de mots). Et j'ai une liste imbriquée de mots-clés définis

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']]

J'avais besoin de vérifier combien de jetons tombent dans chaque liste de mots clés et d'ajouter le résultat en tant que nouvelle colonne de la trame de données existante. Par exemple: si tokens =["become", "farmer","rally","workers","student"] le résultat sera -> [1,2,0]

La fonction suivante a fonctionné comme prévu.

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

J'ai utilisé udf sous withColumn pour accéder à la fonction et j'obtiens une erreur. Je pense qu'il s'agit de passer une liste externe à un udf. Existe-t-il un moyen de passer la liste externe et la colonne datafram à un udf et d'ajouter une nouvelle colonne à mon dataframe?

topicWord = udf(label_maker_topic,StringType())
myDF=myDF.withColumn("topic_Word_count",topicWord(myDF.bodyText_token,keyword_list))
15
Jay

La solution la plus propre consiste à passer des arguments supplémentaires à l'aide de la fermeture:

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())

Cela ne nécessite aucune modification dans keyword_list ou la fonction que vous encapsulez avec UDF. Vous pouvez également utiliser cette méthode pour passer un objet arbitraire. Cela peut être utilisé pour passer par exemple une liste de sets pour des recherches efficaces.

Si vous souhaitez utiliser votre UDF actuel et passer topic_words directement, vous devrez d'abord le convertir en un littéral de colonne:

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()

En fonction de vos données et de vos besoins, il existe des solutions alternatives plus efficaces, qui ne nécessitent pas de FDU (exploser + agréger + réduire) ou de recherches (hachage + opérations vectorielles).

30
zero323

Ce qui suit fonctionne très bien où tout paramètre externe peut être passé à l'UDF (un code modifié pour aider n'importe qui)

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