統計 Pyspark 出現條件時兩個事件之間的月數

Question

我在 Pyspark 工作，當滿足條件時，我需要計算兩個事件之間的月數。

接下來我展示一下我的桌子是怎樣的，以便您更好地了解我。 這是我最初的 dataframe。

from pyspark.sql import Row, Window
from pyspark.sql.functions import *
from datetime import datetime, date

row = Row("id", "start", "condition")
df = sc.parallelize([
    row(1, "2015-01-31", 0),
    row(1, "2015-02-28", 0),
    row(1, "2015-03-31", 0),
    row(1, "2015-04-30", 0),
    row(1, "2015-05-31", 1),
    row(1, "2015-06-30", 1)
]).toDF().withColumn("start", col("start").cast("date"))

## +---+----------+----------+
## | id|     start| condition| 
## +---+----------+----------+
## |  1|2015-01-31|         0|
## |  1|2015-02-28|         0|
## |  1|2015-03-31|         0|
## |  1|2015-04-30|         0|
## |  1|2015-05-31|         1|
## |  1|2015-06-30|         1|
## +---+----------+----------+

我想要這個結果：

## +---+----------+----------+------------------+       
## | id|     start| condition| Months_between|
## +---+----------+----------+---------------+
## |  1|2015-01-31|         0|              4|
## |  1|2015-02-28|         0|              3|
## |  1|2015-03-31|         0|              2|
## |  1|2015-04-30|         0|              1|
## |  1|2015-05-31|         1|              0|
## |  1|2015-06-30|         1|              0|
## +---+----------+----------+---------------+

我想知道條件從 0 變為 1 的一行和另一行之間經過了多少個月。如果條件從未變為 1，則應該為 0。樣本有一個 id，但每個都有很多 id日期。

我曾考慮過制作 window 但我不知道如何獲得月數。 我曾想過這樣的事情：

max_days = (df.select(max("start")).collect()[0][0] - df.select(min("start")).collect()[0][0]).days
days = lambda i: i * 86400
window = Window.partitionBy("id").orderBy(col("start").cast("long")).rangeBetween(days(max_days), 0)

謝謝：它有效：

df2 = df.withColumn(
    'Months_between', 
    F.when(
        F.col('condition') == 0, 
        F.months_between(
            F.min(
                F.when(F.col('condition') == 1, F.col('start'))
            ).over(Window.partitionBy('id')), 
            F.col('start')
        ).cast('int')
    ).otherwise(0)
)

但是當我有這個例子時我發現了一個問題。 當值在 0 和 1 之間多次變化時。

| id|     start|condition|Months_between|
+---+----------+---------+------------------+
|  1|2015-01-31|        0|              2|
|  1|2015-02-28|        0|              1|
|  1|2015-03-31|        1|              0|
|  1|2015-04-30|        1|              0|
|  1|2015-05-31|        0|             -1|
|  1|2015-06-30|        1|              0|
+---+----------+---------+-----------------+

在 2015-05-31 日期，它應該取值 1，但在搜索最小值時，它會得到值 -1。 有什么建議嗎？ 謝謝！

謝謝你的幫助！

Answer 1

您可以找到每個 id 的 condition = 1 的最早日期，並使用months_between獲取該日期與當前行之間的月份：

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'next_start',
    F.first(
        F.when(F.col('condition') == 1, F.col('start')),
        ignorenulls=True
    ).over(
        Window.partitionBy('id')
              .orderBy('start')
              .rowsBetween(0, Window.unboundedFollowing)
    )
).withColumn(
    'Months_between', 
    F.when(
        F.col('condition') == 0, 
        F.months_between(
            F.col('next_start'), 
            F.col('start')
        ).cast('int')
    ).otherwise(0)
).drop('next_start')

df2.show() 
+---+----------+---------+--------------+
| id|     start|condition|Months_between|
+---+----------+---------+--------------+
|  1|2015-01-31|        0|             2|
|  1|2015-02-28|        0|             1|
|  1|2015-03-31|        1|             0|
|  1|2015-04-30|        1|             0|
|  1|2015-05-31|        0|             1|
|  1|2015-06-30|        1|             0|
+---+----------+---------+--------------+