根据 PySpark 中的重复列行创建两个数组

Question

我正在使用 PySpark 并有一个数据框，该数据框有两列

a

和

b

，开始时数据框中的每列/行只有一个值。跨行的

b

中可能（但并不总是）存在重复值。目标是将数据帧转换为在单行中包含

b

中所有重叠值的数组，以及

a

对应值的数组，并在此过程中删除旧行。

b

不与另一行重叠的行应保持不变。

如果这是我的数据：

data = [
 ('00003-01', 4249300705),
 ('00003-01', 4242100870),
 ('00004-10', 4242100870),
 ('00004-10', 4242180791),
 ('00005-01', 4249301111),
 ('00005-01', 4242184444),
 ('00006-10', 4242184444)
]

columns = ["a", "b"]

df = spark.createDataFrame(data, schema=columns)

运行该方法后，

df.show(truncate=False)

将是：

+-----------------------+------------------------------------------------
|a                      |b                                               |
+------------------------------------------------------------------------
|[00003-01, 00004-10]   |[4249300705, 4242180791, 4242100870]            |
|[00005-01, 00006-10]   |[4249301111, 4242184444]                        |
+-----------------------+------------------------------------------------

我已经尝试了几种使用连接的方法，但似乎当有两个以上的实例需要将它们网格在一起形成一行时（不排除这一点 - 我只是没有运气），这种方法不工作得很好。我认为这可以通过一些嵌套循环来完成（如果这是唯一的方法，我并不“完全”反对），但我觉得必须有更聪明和更有效的方法来解决这个问题。

Answer 1

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list, array_distinct, col

spark = SparkSession.builder.master("local").appName("Overlap Rows").getOrCreate()
data = [
    ('00003-01', 4249300705),
    ('00003-01', 4242100870),
    ('00004-10', 4242100870),
    ('00004-10', 4242180791),
    ('00005-01', 4249301111),
    ('00005-01', 4242184444),
    ('00006-10', 4242184444)
]
columns = ["a", "b"]
df = spark.createDataFrame(data, schema=columns)

df_grouped = df.groupBy("b").agg(array_distinct(collect_list("a")).alias("a"))
df_transformed = df_grouped.groupBy("a").agg(collect_list("b").alias("b"))

df_result = df_transformed.select(
    array_distinct(col("a")).alias("a"),
    array_distinct(col("b")).alias("b")
)

df_result.show(truncate=False)

Answer 2

df
.groupBy("a")
.agg(collect_list("b").alias("b"))
.withColumn("b", expr("ARRAY_DISTINCT(FLATTEN(FILTER(COLLECT_LIST(b) OVER(ORDER BY 1), e -> e == b OR ARRAYS_OVERLAP(e, b))))"))
.groupBy("b")
.agg(collect_list("a").as("a"))
.select("a", "b")
.show(5, False)

+--------------------+------------------------------------+
|a                   |b                                   |
+--------------------+------------------------------------+
|[00003-01, 00004-10]|[4249300705, 4242100870, 4242180791]|
|[00005-01, 00006-10]|[4249301111, 4242184444]            |
+--------------------+------------------------------------+