假设我有一个包含三列的表格:
dt
、id
和value
。
df_tmp = spark.createDataFrame([('2023-01-01', 1001, 5),
('2023-01-15', 1001, 3),
('2023-02-10', 1001, 1),
('2023-02-20', 1001, 2),
('2023-01-02', 1002, 7),
('2023-01-02', 1002, 6),
('2023-01-03', 1002, 1)],
["date", "id", "value"])
df.show()
# +----------+----+-----+
# | date| id|value|
# +----------+----+-----+
# |2023-01-01|1001| 5|
# |2023-01-15|1001| 3|
# |2023-02-10|1001| 1|
# |2023-02-20|1001| 2|
# |2023-01-02|1002| 7|
# |2023-01-02|1002| 6|
# |2023-01-03|1002| 1|
# +----------+----+-----+
我想计算每个
value
按id
分组的date
的30天滚动总和,此外,还计算id
被看到的不同日期。看起来像这样的东西:
+----------+----+-----+----------------+-------------------------+
| date| id|value|30_day_value_sum|days_seen_in_past_30_days|
+----------+----+-----+----------------+-------------------------+
|2023-01-01|1001| 5| 0| 0|
|2023-01-15|1001| 3| 0| 1|
|2023-02-10|1001| 1| 3| 1|
|2023-02-20|1001| 2| 1| 2|
|2023-01-02|1002| 7| 0| 0|
|2023-01-02|1002| 6| 7| 1|
|2023-01-03|1002| 1| 13| 2|
+----------+----+-----+----------------+-------------------------+
我怀疑可以使用
Window
来做到这一点,但不清楚具体的细节。
我假设最初您的数据框中有 date 数据类型,所以我使用了比您的输入稍微修改过的输入:
from pyspark.sql import functions as F, Window as W
df_tmp = spark.createDataFrame(
[('2023-01-01', 1001, 5),
('2023-01-15', 1001, 3),
('2023-02-10', 1001, 1),
('2023-02-20', 1001, 2),
('2023-01-01', 1002, 7),
('2023-01-02', 1002, 6),
('2023-01-03', 1002, 1)],
["date", "id", "value"]
).withColumn('date', F.col('date').cast('date'))
在这种情况下,以下窗口应该可以工作:
w = W.partitionBy('id').orderBy(F.expr("unix_date(date)")).rangeBetween(-30, -1)
df = df_tmp.withColumn('30_day_value_sum', F.sum('value').over(w)) \
.withColumn('days_seen_in_past_30_days', F.count('id').over(w))
df = df.fillna(0, subset=['30_day_value_sum'])
df.show()
# +----------+----+-----+----------------+-------------------------+
# | date| id|value|30_day_value_sum|days_seen_in_past_30_days|
# +----------+----+-----+----------------+-------------------------+
# |2023-01-01|1001| 5| 0| 0|
# |2023-01-15|1001| 3| 5| 1|
# |2023-02-10|1001| 1| 3| 1|
# |2023-02-20|1001| 2| 1| 1|
# |2023-01-01|1002| 7| 0| 0|
# |2023-01-02|1002| 6| 7| 1|
# |2023-01-03|1002| 1| 13| 2|
# +----------+----+-----+----------------+-------------------------+
如果您在“日期”列中有 string 类型,那么您应该使用以下 .orderBy` 子句:
.orderBy(F.expr("unix_date(to_date(date))")).rangeBetween(-30, -1)
更多选项和详细信息位于此处。