基於當前行創建新列，計算涉及其他行

Question

假設我有這樣的 dataframe ：

id,month,price
1,2021-04-31,9
1,2021-01-31,5
1,2021-02-31,6
1,2021-03-31,8

因此，對於每個相同的ID ，我想獲取當前行month-1和-2的price總和

例如，對於第1,march,8行，我將在新列中獲得jan為5+6=11因為對於當前 3 月行過去兩個月是 1 月和feb

主數據中也會有其他ids

Answer 1

將month名稱轉換為月份編號，然后將其用於在按id分區的 Window 中進行排序以運行總和：

from pyspark.sql import functions as F, Window

df = spark.createDataFrame([
    (1, "apr", 9), (1, "jan", 5),
    (1, "feb", 6), (1, "march", 8)
], ["id", "month", "price"])

# handle both full and short textual representation of month names
month_number = F.when(F.length("month") == 3, F.month(F.to_date(F.col("month"), "MMM"))) \
    .otherwise(F.month(F.to_date(F.col("month"), "MMMM")))

w = Window.partitionBy("id").orderBy(month_number).rangeBetween(-2, -1)

df.withColumn("price_sum", F.sum("price").over(w)).show()

#+---+-----+-----+---------+
#| id|month|price|price_sum|
#+---+-----+-----+---------+
#|  1|  jan|    5|     null|
#|  1|  feb|    6|        5|
#|  1|march|    8|       11|
#|  1|  apr|    9|       14|
#+---+-----+-----+---------+

---

對於您更新的問題，您可以將日期截斷為月份單位，然后使用 window ，范圍介於interval -2 months和interval -1 months之間：

df = spark.createDataFrame([
    (1, "2021-04-30", 9), (1, "2021-01-31", 5),
    (1, "2021-02-28", 6), (1, "2021-03-31", 8)
], ["id", "month", "price"])

df.withColumn(
    "date",
    F.date_trunc("month", F.col("month"))
).withColumn(
    "price_sum",
    F.expr("""sum(price) over(partition by id order by date 
                              range between interval 2 months preceding 
                              and interval 1 months preceding)
    """)
).drop("date").show()

Answer 2

使用 window function 對價格、按月分區（數字）求和，並使用 ROWS 獲得前兩行。

from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()

df = spark.read.option("header",True).csv("path/to/file") # Assuming file is csv
df.createOrReplaceTempView('df')
df1 = spark.sql("""
           SELECT id,month,price,
             CASE
               WHEN month = 'jan' THEN 1
               WHEN month = 'feb' THEN 2
               .
               .
               .
               else 12
             END AS month_num
           FROM df 
           """)

df1.createOrReplaceTempView('df1')
spark.sql("""
           SELECT id, month, price,
           SUM(price) OVER (PARTITION BY id ORDER BY month_num ROWS 2 PRECEDING) AS price_sum
           FROM df1 
           """).show()

添加到第二個查詢

WHERE month_num NOT IN (1,2)

如果您想要 1 月和 2 月的 price_sum 0。