将数据框列和外部列表传递给withColumn下的udf

时间:2016-05-24 09:42:34

标签: python apache-spark pyspark apache-spark-sql user-defined-functions

我有一个具有以下结构的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']]

我需要检查每个关键字列表下有多少令牌,并将结果添加为现有数据帧的新列。 例如:if 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

我在withColumn下使用了udf来访问该函数,但是我收到了一个错误。我认为这是关于将外部列表传递给udf的。有没有办法可以将外部列表和datafram列传递给udf并向我的数据框添加新列?

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

4 个答案:

答案 0 :(得分:29)

最干净的解决方案是使用闭包传递其他参数:

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(爆炸+聚合+折叠)或查找(散列+矢量操作)。

答案 1 :(得分:8)

以下工作正常,可以将任何外部参数传递给UDF(一个可以帮助任何人的调整代码)

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

答案 2 :(得分:0)

以另一种方式使用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()

答案 3 :(得分:0)

如果列表很大,应将keyword_list列表广播到群集中的所有节点。我猜零的解决方案行得通,因为该列表很小,并且会自动广播。我认为最好进行明确广播,以免产生疑问(对于较大的列表,明确广播是必需的。)

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

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

    return twt_list

def make_topic_word_better(topic_words_broadcasted):
    def f(c):
        return label_maker_topic(c, topic_words_broadcasted)
    return F.udf(f)

df = spark.createDataFrame([["union",], ["party",]]).toDF("tokens")
b = spark.sparkContext.broadcast(keyword_list)
df.withColumn("topics", make_topic_word_better(b)(F.col("tokens"))).show()

以下是输出内容:

+------+---------+
|tokens|   topics|
+------+---------+
| union|[0, 0, 0]|
| party|[0, 0, 0]|
+------+---------+

请注意,您需要致电value来访问已广播的列表(例如topic_words_broadcasted.value)。有关如何在PySpark分析中广播列表和词典的更多详细信息,请参见this blog post。这是一个困难的实现,但是掌握起来很重要,因为许多PySpark UDF依赖于广播的列表或词典。