PySpark - 将单个整数列表与列表列进行比较

时间:2018-01-31 17:08:30

标签: python apache-spark pyspark

我试图检查spark数据框(带列表的列)中的哪些条目包含给定列表中最大数量的值。

我提出的最佳方法是使用rdd.foreach()迭代数据框,并使用python' s set1.intersection(set2)将给定列表与每个条目进行比较。

我的问题是,spark是否具有任何内置功能,因此可以避免使用.foreach进行迭代?

感谢您的帮助!

P.S。我的数据框看起来像这样:

+-------------+---------------------+                                           
|   cardnumber|collect_list(article)|
+-------------+---------------------+
|2310000000855| [12480, 49627, 80...|
|2310000008455| [35531, 22564, 15...|
|2310000011462| [117112, 156087, ...|
+-------------+---------------------+

我试图在第二列中找到包含最多交叉点的条目,其中包含给定的文章列表,例如[151574, 87239, 117908, 162475, 48599]

2 个答案:

答案 0 :(得分:1)

您可以在数据帧中尝试相同的设置操作,而不是使用rdd.foreach:

my_udf=udf(lambda A,B: list(set(A).intersection(set(B))))
df=df.withColumn('intersect_value', my_udf('A', 'B'))

您可以使用len函数在UDF本身中获取交叉列表的大小,并从此数据帧中执行所需的操作。

答案 1 :(得分:1)

这里唯一的替代方案是udf,但它没有多大区别。

from pyspark.sql.functions import udf, li, col

def intersect(xs):
    xs = set(xs)
    @udf("array<long>")
    def _(ys):
        return list(xs.intersection(ys))
    return _

它可以应用为:

a_list = [1, 4, 6]

df = spark.createDataFrame([
    (1, [3, 4, 8]), (2, [7, 2, 6])
], ("id", "articles"))

df.withColumn("intersect", intersect(a_list)("articles")).show()

# +---+---------+---------+
# | id| articles|intersect|
# +---+---------+---------+
# |  1|[3, 4, 8]|      [4]|
# |  2|[7, 2, 6]|      [6]|
# +---+---------+---------+

根据名称,您似乎使用collect_list,因此您的数据看起来可能是这样的:

df_long = spark.createDataFrame([
    (1, 3),(1, 4), (1, 8), (2, 7), (2, 7), (2, 6)
], ("id", "articles"))

在这种情况下问题更简单。加入

lookup = spark.createDataFrame(a_list, "long").toDF("articles")

joined = lookup.join(df_long, ["articles"])

并汇总结果:

joined.groupBy("id").count().show()
# +---+-----+                                                                     
# | id|count|
# +---+-----+
# |  1|    1|
# |  2|    1|
# +---+-----+


joined.groupBy("id").agg(collect_list("articles")).show()
# +---+----------------------+                                                    
# | id|collect_list(articles)|
# +---+----------------------+
# |  1|                   [4]|
# |  2|                   [6]|
# +---+----------------------+