请注意,下面不使用编码,但它使用连接完全符合您的要求。一起来看看吧。
必要的进口
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window
创建nodes_df
val data1 = Seq(("AAA",1), ("BBB",1), ("CCC",2))
val nodes_df = data1.toDF("node", "group")
nodes_df.show()
+----+-----+
|node|group|
+----+-----+
| AAA| 1|
| BBB| 1|
| CCC| 2|
+----+-----+
创建节点
使用row_number 和Window 添加行号。我将 orderBy 从 group 更改为 node,因为您无法保证 AAA 是否为 1,BBB 是否为 2,反之亦然,因为它们具有相同的 group 编号。
val nodes = nodes_df.withColumn("id",row_number().over(Window.orderBy("node")))
nodes.show()
+----+-----+---+
|node|group| id|
+----+-----+---+
| AAA| 1| 1|
| BBB| 1| 2|
| CCC| 2| 3|
+----+-----+---+
创建 arcs_df
val data2 = Seq(("AAA","BBB",1), ("BBB","CCC",1))
val arcs_df = data2.toDF("node_from", "node_to", "weight")
arcs_df.show()
+---------+-------+------+
|node_from|node_to|weight|
+---------+-------+------+
| AAA| BBB| 1|
| BBB| CCC| 1|
+---------+-------+------+
执行连接以获得所需的输出
现在我们基本上是您想知道如何将id 从nodes 映射到arcs_df 中的从/到节点的地方。我通过使用两个连接来实现这一点。首先是从nodes 到arcs_df,创建一个中间数据帧first_join_df,用于说明目的。然后从nodes 到first_join_df。在加入之后,我删除了我不关心的列,例如group,并将id 列重命名为node_from/node_to(如果适用)。我以.select(... 结束,以获得您要求的订单。看看,让我知道是否有任何不清楚的地方。希望这会有所帮助!
val first_join_df = arcs_df.join(nodes, arcs_df.col("node_from") === nodes.col("node"), "leftouter")
.drop("node_from", "node", "group")
.withColumnRenamed("id", "node_from")
first_join_df.show()
+-------+------+---------+
|node_to|weight|node_from|
+-------+------+---------+
| BBB| 1| 1|
| CCC| 1| 2|
+-------+------+---------+
val second_join_df = first_join_df.join(nodes, first_join_df.col("node_to") === nodes.col("node"), "leftouter")
.drop("node_to", "node", "group")
.withColumnRenamed("id", "node_to")
.select("node_from", "node_to", "weight")
second_join_df.show()
+---------+-------+------+
|node_from|node_to|weight|
+---------+-------+------+
| 1| 2| 1|
| 2| 3| 1|
+---------+-------+------+