Pyspark - «Рекурсивная» функция с участием последнего дня

Question

Я работаю над процессом в pyspark, в котором у меня есть фрейм данных, и пытаюсь добавить еще один столбец (используя метод withColumn).

Проблема в том, что формула:

STATUS1 = Если «PETP-today»> 0, то «Status1 last day» + «PETP-today», иначе 0

Каждый результат для Status1 включает статус1 из результата последнего дня.

Одно решение, которое я нашел, состояло в том, чтобы создать фрейм данных pandas и запускать записи по одной, пока я не смогу вычислить каждый из них, используя переменные. Однако у меня будут проблемы с производительностью. Вы можете помочь?

Рассмотрим столбцы данных: дата (ежедневно) / PETP (число с плавающей запятой) / STATUS1? (Поплавок)

Я действительно ценю любую помощь!

Answer 1

Я думаю, что ключом к вашему решению является функция lag . Попробуйте это (для простоты я предполагаю целочисленные данные для всех столбцов):

Сначала сдвиньте столбец на один день вверх

import pyspark
from pyspark.sql import SparkSession
from pyspark import SparkContext
import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql import Window

sc = SparkContext.getOrCreate()
spark = SparkSession(sc)

columns = ['date', 'petp', 'status']
data = [(0, 0, 0), (1, 1, 1), (2, 2, 2), (3,3,3), (4,4,4), (5,5,5)]
pd_data = pd.DataFrame.from_records(data=data, columns=columns)
spark_data = spark.createDataFrame(pd_data)

spark_data_with_lag = spark_data.withColumn("status_last_day", F.lag("status", 1, 0).over(Window.orderBy("date")))
spark_data_with_lag.show()

+----+----+------+---------------+
|date|petp|status|status_last_day|
+----+----+------+---------------+
|   1|   1|     1|              0|
|   2|   2|     2|              1|
|   3|   3|     3|              2|
|   4|   4|     4|              3|
|   5|   5|     5|              4|
+----+----+------+---------------+

Затем используйте эти данные в вашем условном

status2 = spark_data_with_lag.withColumn("status2", F.when(F.col("date") > 0, F.col("petp") + F.col("status_last_day")).otherwise(0))
status2.show()

+----+----+------+---------------+-------+
|date|petp|status|status_last_day|status2|
+----+----+------+---------------+-------+
|   1|   1|     1|              0|      1|
|   2|   2|     2|              1|      3|
|   3|   3|     3|              2|      5|
|   4|   4|     4|              3|      7|
|   5|   5|     5|              4|      9|
+----+----+------+---------------+-------+

Надеюсь, это то, что вы искали.