Pyspark:为每个不同的值向行添加增量计数器

Question

我需要向数据帧添加一个"行号"，但对于列中的每个新值，必须重新启动此"行号"。

我举个例子：

from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('test').getOrCreate()
df = spark.createDataFrame([
    ('2018-01-01', 'John', 0),
    ('2018-01-01', 'Paul', 1),
    ('2018-01-08', 'Paul', 3),
    ('2018-01-08', 'Pete', 4),
    ('2018-01-08', 'John', 3),
    ('2018-01-15', 'Mary', 6),
    ('2018-01-15', 'Pete', 6),
    ('2018-01-15', 'John', 6),
    ('2018-01-15', 'Paul', 6),
], ['str_date', 'name', 'value'])
# Convert str_date to date:
df = df.withColumn('date', to_date(df['str_date'])) 
    .select(['date', 'name', 'value'])
# Sort by name and date
df.orderBy(['name', 'date']).show()
## +----------+----+-----+
## |      date|name|value|
## +----------+----+-----+
## |2018-01-01|John|    0|
## |2018-01-08|John|    3|
## |2018-01-15|John|    6|
## |2018-01-15|Mary|    6|
## |2018-01-01|Paul|    1|
## |2018-01-08|Paul|    3|
## |2018-01-15|Paul|    6|
## |2018-01-08|Pete|    4|
## |2018-01-15|Pete|    6|
## +----------+----+-----+

因此，我需要的是添加一个新列，其中包含每个name的行号：

# Expected result
## +----------+----+-----+------+
## |      date|name|value|rowNum|
## +----------+----+-----+------+
## |2018-01-01|John|    0|     1| <- First row for 'John'
## |2018-01-08|John|    3|     2|
## |2018-01-15|John|    6|     3|
## |2018-01-15|Mary|    6|     1| <- First row for 'Mary'
## |2018-01-01|Paul|    1|     1| <- First row for 'Paul'
## |2018-01-08|Paul|    3|     2|
## |2018-01-15|Paul|    6|     3|
## |2018-01-08|Pete|    4|     1| <- First row for 'Pete'
## |2018-01-15|Pete|    6|     2|
## +----------+----+-----+------+

我一直在尝试使用 Window 函数，但我卡住了。你能帮帮我吗？

笔记：

保证对行
• 进行排序（如果未排序，则将作为工作管道的一部分对行进行排序）
• 我正在使用 Spark 2.4.0

Answer 1

使用像row_number这样的排名函数来执行此操作。如果某个名字在给定日期可以有联系，请改用dense_rank。

from pyspark.sql import Window 
from pyspark.sql import functions as f
#Window definition
w = Window.partitionBy(df.name).orderBy(df.date)
res = df.withColumn('rnum',f.row_number().over(w))
res.show()

Answer 2

瓦姆西的回答是正确的。 错过了一个（）row_number所以...

    w = Window.partitionBy(df.name).orderBy(df.date)
    res = df.withColumn('rnum',f.row_number().over(w))  # change after row_number
    res.show()