python - 使用搜索和条件查找两列值之间的差异

Question

在 pyspark 中，我有一个如下所示的数据框，其中的行根据 id 和 k1 的值进行排序。此外，每一行都有一个分配给它的唯一升序号 (rowid)。

-----------------------
rowid | id | k1  | k2 |
-----------------------
1     | 1  | v1 | l1  |
2     | 1  | v1 | v1  |
3     | 1  | v1 | l2  |
4     | 2  | v2 | v2  |
5     | 2  | v2 | l3  |
6     | 3  | v3 | l3  |
----------------------

对于id的每个唯一值，我想计算k1==k2的第一行的rowid与观察到该id的记录的第一行对应的rowid + 1之间的差值，并且将结果存储在新列中(即排名)。输出应如下所示。

----------------
 id | k1  |rank |
-----------------
 1  | v1  | 2   |
 2  | v2  | 1   |
 3  | v3  | 0   | 
-----------------

例如，对于id=1，当rowid=2时k1==k2的值。第一次观察到id=1是在rowid=1的时候。将 2-1+1=2 放在排名列中。对于 id =3，我们没有任何记录中列 k1 和 k2 的值匹配。因此，用 0(或空值)填充排名列。

我假设这涉及一个基于 id 的 groupBy，但我不确定如何获取与列 k1 和 k2 匹配的行对应的索引以及对应于每个唯一 id 的第一个 rowid。

Answer 1

您可以在 id 和 k1 上使用带有 groupBy 的 API 函数来执行此操作，这应该是 faster than using a udf :

import pyspark.sql.functions as f

df.groupBy("id", "k1")\
    .agg(
        f.min(f.when(f.col("k1")==f.col("k2"), f.col("rowid"))).alias("first_equal"),
        f.min("rowid").alias("first_row")
    )\
    .select("id", "k1", (f.col("first_equal")-f.col("first_row")+1).alias("rank"))\
    .fillna(0)\
    .show()
#+---+---+----+
#| id| k1|rank|
#+---+---+----+
#|  1| v1|   2|
#|  2| v2|   1|
#|  3| v3|   0|
#+---+---+----+

rank 的计算可以分解为两个聚合步骤:

• 第一个聚合采用最小 rowid，对于每个 id、k1 对，k1==k2 .
• 第二个聚合采用每个 id、k1 对的最小 rowid。

您取这些差值(+1 根据您的要求)，最后用 0 填充任何 null 值。

---

更新:另一种使用row_number的方法:

from pyspark.sql import Window

# you can define your own order by column
w = Window.partitionBy("id", "k1").orderBy("rowid")

df.withColumn("rank", f.when(f.expr("k1 = k2"), f.row_number().over(w)))\
    .groupBy("id", "k1")\
    .agg(f.min("rank"))\
    .fillna(0)\
    .show()
# Same as above