PySpark窗口函数：按var1统计近3天内相同值的行数

嘿，我来帮你搞定这个PySpark的统计需求！要实现每一行数据统计过去3天内相同var1值的记录行数，用窗口函数就能轻松解决，具体步骤和代码如下：

1. 准备数据并处理日期类型 #

首先得把字符串格式的date转成PySpark的日期类型，这是后续窗口计算的基础。另外你的测试数据最后一条不完整，我先帮你补全了一条合理的记录：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
from pyspark.sql.window import Window
from pyspark.sql.functions import col, count, to_date

# 初始化SparkSession
spark = SparkSession.builder.appName("Past3DaysRecordCount").getOrCreate()

# 补全后的测试数据
_schema = StructType([
    StructField("date", StringType(), True),
    StructField("var1", IntegerType(), True),
    StructField("var2", StringType(), True)
])
test_list = [
    ('2017-01-30',123,'A'), 
    ('2017-01-17',123,'B'), 
    ('2017-01-15',123,'A'), 
    ('2017-01-15',123,'A'), 
    ('2017-01-14',123,'A'), 
    ('2017-01-11',123,'B'), 
    ('2017-01-29',456,'A'), 
    ('2017-01-22',789,'B'), 
    ('2017-01-28',789,'B')
]

# 创建DataFrame并转换日期格式
df = spark.createDataFrame(test_list, schema=_schema)
df = df.withColumn("date", to_date(col("date"), "yyyy-MM-dd"))

2. 定义关键的窗口规范 #

这里核心是用范围窗口（Range Window），因为我们要统计的是时间区间内的记录，而不是固定行数（毕竟日期可能不连续）。窗口要满足三个条件：

• 按var1分区：只统计相同var1组内的记录
• 按date排序：确保时间顺序正确
• 范围是当前日期往前推3天到当前日期：这里把日期转成时间戳（秒数），用rangeBetween指定区间，3天就是3*86400秒

# 定义窗口规则：按var1分组，按日期时间戳排序，范围是过去3天到当天
window_spec = Window.partitionBy("var1")\
                   .orderBy(col("date").cast("long"))\
                   .rangeBetween(-3*86400, 0)

3. 计算过去3天的记录数 #

用count函数在窗口内统计行数，就能得到每一行对应的过去3天同var1的记录数了：

# 计算并添加统计列
result_df = df.withColumn("past_3d_record_count", count("var1").over(window_spec))

# 查看排序后的结果
result_df.orderBy("var1", "date").show()

结果解释 #

举几个例子帮你理解：

• 对于2017-01-30、var1=123的记录，过去3天是2017-01-27到2017-01-30，这个区间只有它自己，所以统计数是1
• 对于2017-01-15的两条var1=123的记录，过去3天是2017-01-12到2017-01-15，包含2017-01-14的1条和当天的2条，所以统计数是3
• 对于2017-01-28、var1=789的记录，过去3天是2017-01-25到2017-01-28，只有它自己，而2017-01-22的那条超出了3天范围，所以统计数是1

可选调整 #

如果你不想把当前行计入统计，只算过去3天的历史记录，只要把窗口范围改成rangeBetween(-3*86400, -1)就可以了。

内容的提问来源于stack exchange，提问作者B_Miner