python - 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 pls中提供逻辑帮助。我真的需要克服这个障碍，请提供代码帮助

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|
+---+---------+