从pyspark datafame创建NXN矩阵

时间:2020-03-20 17:50:25

标签: pyspark pyspark-dataframes

我对pyspark很陌生。我有1万个文本数据集。我使用Minhash lsh创建一个Jaccard距离。 我得到的输出例如

col1 col2 dist
A   B       0.77
B   C       0.56
C   A       0.88

我想将其转换为NxN矩阵格式。

    A        B      C
A   0       0.77    0.88

B   0.77     0      0.56

C   0.88    0.56    0

有什么方法可以使用pyspark创建它。 我很感谢这些建议。

1 个答案:

答案 0 :(得分:1)

可以使用下面的代码来完成。但是,由于groupBypivotsunion,然后是groupBy,因此它将非常耗费计算资源。之所以有两个 groupBy枢纽,是因为数据A-BB-A中有两个组合。

df1=df.groupBy("col1").pivot("col2").agg(F.first("dist")).orderBy("col1")
df2=df.groupBy(F.col("col2").alias("col1")).pivot("col1").agg(F.first("dist")).orderBy("col1")
df3=df1.union(df2)
df3.groupBy("col1")\
.agg(*(F.first(x,True).alias(x) for x in df3.columns if x != 'col1'))\
.fillna(0)\
.orderBy("col1")\
.show()

+----+----+----+----+
|col1|   A|   B|   C|
+----+----+----+----+
|   A| 0.0|0.77|0.88|
|   B|0.77| 0.0|0.56|
|   C|0.88|0.56| 0.0|
+----+----+----+----+