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()