按交集对 pyspark 数据框进行分组

Question

我需要按列中数组的交集对 PySpark 数据帧进行分组。例如来自这样的数据框：

v1 | [1, 2, 3]
v2 | [4, 5]
v3 | [1, 7]

结果应该是：

[v1, v3] | [1, 2, 3, 7]
[v2] | [4, 5]

因为第 1 行和第 3 行的值都是 1。

有交集时group by之类的方法吗？

提前感谢您提供解决此问题的想法和建议。

Answer 1

from pyspark.sql import functions as F

df = spark.createDataFrame([["v1", [1,2,3]], ["v2", [4,5]], ["v3",[1,7]]],["id","arr"])

df1= df.select("*", F.explode("arr").alias("explode_arr")).groupBy("explode_arr").agg(F.collect_set("id").alias("ids"))

df2=df.select("*", F.explode("arr").alias("explode_arr")).join(df1, ["explode_arr"],\
    "inner").groupBy("ids").agg(F.collect_set("arr").alias("array_set")).\
    select("ids",F.array_distinct(F.expr("flatten(array_set)")).alias("intersection_arrays"))

df3= df2.where(F.size("ids")>1).select(F.explode("ids").alias("ids")).select(F.array("ids").alias("ids"))

df4= df2.join(df3.withColumn("flag", F.lit(1)),["ids"],"left_outer").where(F.col("flag").isNull()).drop("flag")

df4.show()

+--------+-------------------+
|     ids|intersection_arrays|
+--------+-------------------+
|    [v2]|             [4, 5]|
|[v3, v1]|       [1, 7, 2, 3]|
+--------+-------------------+

按交集对 pyspark 数据框进行分组

Grouping pyspark dataframe by intersection

grouping

intersection

dataframe

apache-spark

pyspark