Question

在pyspark中，我有一个像下面这样的数据框，其中的行根据id和k1的值排序。此外，每行都有一个唯一的升序编号（行）。

-----------------------
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的记录+ 1的第一行的rowid之间的差，以及将结果存储在新列（即等级）中。输出应如下所示。

----------------
 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（或null）填充等级列。

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

Answer 1

首先创建一个示例数据框，

import pyspark.sql.functions as F
from pyspark.sql.types import *

df = sql.createDataFrame([
            (1, 1, 'v1' , 'l1'),
            (2, 1, 'v1' , 'v1'),
            (3, 1, 'v1' , 'l2'),
            (4, 2, 'v2' , 'v2'),
            (5, 2, 'v2' , 'l3'),
            (6, 3, 'v3' , 'l3'),
            ],[
            'rowid', 'id', 'k1', 'k2'])

然后创建一个udf并将其应用于列，

def get_rank_udf(rows):
    rows = sorted(rows, key=lambda x: x['rowid'])
    first_row_id = rows[0]['rowid']
    for _r in rows:
        if _r['k1'] == _r['k2']:
            equal_row_id = _r['rowid']
            break
        else:
            equal_row_id = None

    if equal_row_id is None:
        return 0
    return equal_row_id - first_row_id + 1

get_rank = F.udf(lambda x: get_rank_udf(x), IntegerType())

df = df.groupby('id', 'k1').agg(F.collect_list(F.struct('rowid', 'k1', 'k2')).alias('elements'))\
       .withColumn('rank', get_rank(F.col('elements')))\
       .select('id', 'k1', 'rank')

这给出了输出，

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

Answer 2

您可以使用API函数并在groupBy和id（应为faster than using a udf）上使用k1：

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是每个k1==k2，id对的k1。
第二个聚合在每对rowid和id上占用最小k1。

您将这些差异（根据您的要求{+1）并最终用null填充任何0值。

更新：使用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

通过搜索和条件找出两个列值之间的差异

2 个答案: