pyspark对角线比较两列

Question

问题陈述：在 pyspark 中，我必须比较输入数据框下方的 ex 对角线的两列，我必须比较 stn_fr_cd 和 stn_to_cd ，即对于 val_no 1，有 2 行。现在我必须将第一行的 stn_fr_cd 与第二行的 stn_to_cd 和第一行的 stn_to_cd 与第二行的 stn_fr_cd 进行比较。

从下面的输入数据帧，由于 val_no stn_fr_Cd 和 stn_to_cd 诊断元素都相等，我将我的值增加为 1

下面是我的输入，有 4 列 id,val_no,stn_fr_cd,stn_to_cd

id val_no    stn_fr_cd stn_to_cd

8A   1        CPH      GDN                  

8A   1        GDN      CPH                  

8A   2        GDN      CPH                  

8A   2        CPH      GDN                  

8A   3        CPH      GDN                  

8A   3        GDN      CPH                  

8A   4        CPH      GDN                  

8A   4        GDN      CPH 

下面应该是我的输出

8A 4

我如何得到 4 是 val_no 1,2,3,4 stn_fr_cd 和 stn_to_cd 对角线元素都相等

谁能帮我解决pyspark中的逻辑问题。我真的需要跨过这个障碍，请帮忙写代码

Answer 1

我认为这是你想要的，我可能是错的。让我知道它是否适合您，或者我是否可以更新它。我使用窗口函数来获取两列的领先地位，如果它们都相等，则该分区将获得 1，否则为 0，然后仅按 id 分组并将我的检查列相加。我又添加了 2 行（val_no=5）以表明它们没有被选中，因为它们不满足对角线的两个条件。

df.show()

+---+------+---------+---------+
| id|val_no|stn_fr_cd|stn_to_cd|
+---+------+---------+---------+
| 8A|     1|      CPH|      GDN|
| 8A|     1|      GDN|      CPH|
| 8A|     2|      GDN|      CPH|
| 8A|     2|      CPH|      GDN|
| 8A|     3|      CPH|      GDN|
| 8A|     3|      GDN|      CPH|
| 8A|     4|      CPH|      GDN|
| 8A|     4|      GDN|      CPH|
| 8A|     5|      GDN|      GDN|
| 8A|     5|      CPH|      GDN|
+---+------+---------+---------+

from pyspark.sql import functions as F
from pyspark.sql.window import Window
w=Window().partitionBy("id","val_no").orderBy("val_no")

df.withColumn("fr", F.lead("stn_fr_cd").over(w))\
  .withColumn("to", F.lead("stn_to_cd").over(w))\
  .withColumn("check", F.when((F.col("stn_fr_cd")==F.col("to"))&(F.col("stn_to_cd")==F.col("fr")),F.lit(1)).otherwise(F.lit(0)))\
  .groupBy("id").agg(F.sum("check").alias("diagonals")).show()

+---+---------+
| id|diagonals|
+---+---------+
| 8A|        4|
+---+---------+