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

Vamsi的答案是正确的。错过了（）作为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()

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

2 个答案: