количество нулевых дней в поле строки

Question

У меня есть искровой фрейм данных, как в столбце ввода ниже. В нем есть столбец даты «даты» и столбец типа int «кол-во». Я хотел бы создать новый столбец «daysout», в котором будет разница в днях между текущим значением даты и первой последовательной датой, где qty = 0. Ниже я привел пример ввода и вывода. Любые подсказки приветствуются.

входной df:

dates      qty
2020-04-01 1
2020-04-02 0
2020-04-03 0
2020-04-04 3
2020-04-05 0
2020-04-06 7

выход:

dates      qty daysout
2020-04-01 1   0
2020-04-02 0   0
2020-04-03 0   1
2020-04-04 3   2
2020-04-05 0   0
2020-04-06 7   1

Answer 1

Вот возможный подход, который сравнивает, если текущая строка равна 0, а отставленная строка не равна 0, затем берет сумму этого окна, которое затем действует как окно для присвоения номера строки и вычитает 1, чтобы получить желаемое результат:

import pyspark.sql.functions as F

w = Window().partitionBy().orderBy(F.col("dates"))
w1 = F.sum(F.when((F.col("qty")==0)&(F.lag("qty").over(w)!=0),1).otherwise(0)).over(w)
w2 = Window.partitionBy(w1).orderBy('dates')

---

df.withColumn("daysout",F.row_number().over(w2) - 1).show()

+----------+---+-------+
|     dates|qty|daysout|
+----------+---+-------+
|2020-04-01|  1|      0|
|2020-04-02|  0|      0|
|2020-04-03|  0|      1|
|2020-04-04|  3|      2|
|2020-04-05|  0|      0|
|2020-04-06|  7|      1|
+----------+---+-------+