apache-spark - Spark DataFrame：如何在给定条件下为两行分配标签

Question

我有一个数据框如下：

+--------+--------+
| Node 1 | Node 2 |
+--------+--------+
|    A   |    B   |
|    B   |    A   |
|    C   |    D   |
|    D   |    C   |
+--------+--------+

该数据帧保证具有反转的节点对，如 A/B 和 B/A。我想要的是像这样成对标记行：

+--------+--------+-------+
| Node 1 | Node 2 | Label |
+--------+--------+-------+
|    A   |    B   |   1   |
|    B   |    A   |   1   |
|    C   |    D   |   2   |
|    D   |    C   |   2   |
+--------+--------+-------+

在这种情况下，我无法从 A 和 B 中制作标签并对字符串进行排序（.withColumn('Label', alphabetize_string(concat(Node1, Node2))因为它可能会使用从本示例中过滤出来但出现在我的真实数据框中的非倒排节点的标签重复标签。

我该怎么做呢？

Answer 1

使用 PySpark sql 函数 ,least对greatest对进行分组，然后dense_rank获取标签列。

from pyspark.sql.functions import least,greatest
from pyspark.sql import Window 
w = Window.orderBy(least(col('node_1'),col('node_2')),greatest(col('node_1'),col('node_2'))) 
result = df.withColumn('label',dense_rank().over(w))
result.show()