Question

我有一个数据框，如下所示：

---------------
id   | name   |
---------------
 1   | joe    |
 1   | john   |
 2   | jane   |
 3   | jo     |
---------------

目标是，如果“ id”列重复，则从1开始向其添加升序数字。

在Pandas中，我可以这样操作：

count_id = df.groupby(['id']).cumcount()
count_num = count_id.replace(0, '').astype(str)
df['id'] += count_num

我尝试在PySpark中使用相同的逻辑，但没有成功。

结果应为：

id   | name   |
---------------
 1   | joe    |
 11  | john   |
 2   | jane   |
 3   | jo     |
---------------

我如何在PySpark中实现相同目标？任何帮助将不胜感激。

Answer 1

要复制该输出，可以使用Window获取每个id的{{3}}，然后使用row_number将其添加到id

import pyspark.sql.functions as f
from pyspark.sql import Window

w = Window.partitionBy("id").orderBy("name")
df.withColumn("row_number", f.row_number().over(w)-1)\
    .withColumn(
        "id", 
        f.when(
            f.col("row_number") > 0, 
            f.concat(f.col("id"), f.col("row_number"))
        ).otherwise(f.col("id"))
    )\
    .drop("row_number")\
    .show()
#+---+----+
#| id|name|
#+---+----+
#|  1| joe|
#| 11|john|
#|  3|  jo|
#|  2|jane|
#+---+----+

注意：如果尚未将id列转换为StringType列，则将其转换为

。

为了获得您最初在问题中所述的输出作为期望的结果，除了计算行号之外，您还必须concat。仅当计数大于1时，才连接行号。

import pyspark.sql.functions as f
from pyspark.sql import Window

w = Window.partitionBy("id")
df.withColumn("count", f.count("*").over(w))\
    .withColumn("row_number", f.row_number().over(w.orderBy("name")))\
    .withColumn(
        "id", 
        f.when(
            f.col("count") > 1, 
            f.concat(f.col("id"), f.col("row_number"))
        ).otherwise(f.col("id"))
    )\
    .drop("count", "row_number")\
    .show()
#+---+----+
#| id|name|
#+---+----+
#| 11| joe|
#| 12|john|
#|  3|  jo|
#|  2|jane|
#+---+----+

PySpark中的Groupby cumcount

1 个答案: