Как интерполировать временные ряды на основе временного промежутка между ненулевыми значениями в PySpark

Имея этот входной фрейм данных:

df = spark.createDataFrame([
    (1642205833225, 58.00), (1642205888654, float('nan')),
    (1642205899657, float('nan')), (1642205899970, 55.00),
    (1642206338180, float('nan')), (1642206353652, 56.45),
    (1642206853451, float('nan')), (1642207353652, 80.45)
], ["timestamp", "value"])

# replace NaN value by Nulls
df = df.replace(float("nan"), None, ["value"])

Вы можете использовать некоторые оконные функции (last, first), чтобы получить следующие и предыдущие ненулевые значения для каждой строки и рассчитать временной интервал следующим образом:

from pyspark.sql import functions as F, Window

w1 = Window.orderBy("timestamp").rowsBetween(1, Window.unboundedFollowing)
w2 = Window.orderBy("timestamp").rowsBetween(Window.unboundedPreceding, -1)

df = (
    df.withColumn("rn", F.row_number().over(Window.orderBy("timestamp")))
    .withColumn("next_val", F.first("value", ignorenulls=True).over(w1))
    .withColumn("next_rn", F.first(F.when(F.col("value").isNotNull(), F.col("rn")), ignorenulls=True).over(w1))
    .withColumn("prev_val", F.last("value", ignorenulls=True).over(w2))
    .withColumn("prev_rn", F.last(F.when(F.col("value").isNotNull(), F.col("rn")), ignorenulls=True).over(w2))
    .withColumn("timegap_to_next", F.when(F.col("value").isNotNull(), F.min(F.when(F.col("value").isNotNull(), F.col("timestamp"))).over(w1) - F.col("timestamp")))
)

Теперь вы можете выполнить линейную интерполяцию столбца value в зависимости от вашего порога, используя выражение when:

w3 = Window.orderBy("timestamp").rowsBetween(Window.unboundedPreceding, Window.currentRow)

df = df.withColumn(
    "value",
    F.coalesce(
        "value",
        F.when(
            F.last("timegap_to_next", ignorenulls=True).over(w3) < 500000,
            (F.col("prev_val") + 
            ((F.col("next_val") - F.col("prev_val"))/ 
            (F.col("next_timestamp") - F.col("prev_next_timestamp"))
            * (F.col("timestamp") - F.col("prev_next_timestamp")
                    )
                )
            )
        )
    )
).select("timestamp", "value", "timegap_to_next")

df.show()

#+-------------+------+---------------+
#|    timestamp| value|timegap_to_next|
#+-------------+------+---------------+
#|1642205833225|  58.0|          66745|
#|1642205888654|  56.0|           null|
#|1642205899657|  57.0|           null|
#|1642205899970|  55.0|         453682|
#|1642206338180|55.725|           null|
#|1642206353652| 56.45|        1000000|
#|1642206853451|  null|           null|
#|1642207353652| 80.45|           null|
#+-------------+------+---------------+