Есть ли способ добавить несколько столбцов в фрейм данных, рассчитанный на основе скользящих средних из разных столбцов и / или за разную продолжительность

Question

У меня есть фрейм данных с данными временных рядов, и я пытаюсь добавить в него множество столбцов скользящих средних с разными окнами различных диапазонов.Когда я делаю этот столбец за столбцом, результаты довольно медленные.

Я пытался просто накапливать withColumn вызовы до тех пор, пока не получу их все.

Псевдокод:

import pyspark.sql.functions as pysparkSqlFunctions

## working from a data frame with 12 colums:
## - key as a String
## - time as a DateTime
## - col_{1:10} as numeric values

window_1h =  Window.partitionBy("key")                \
                   .orderBy(col("time").cast("long")) \
                   .rangeBetween(-3600, 0)
window_2h =  Window.partitionBy("key")                \
                   .orderBy(col("time").cast("long")) \
                   .rangeBetween(-7200, 0)
df = df.withColumn("col1_1h", pysparkSqlFunctions.avg("col_1").over(window_1h))
df = df.withColumn("col1_2h", pysparkSqlFunctions.avg("col_1").over(window_2h))
df = df.withColumn("col2_1h", pysparkSqlFunctions.avg("col_2").over(window_1h))
df = df.withColumn("col2_2h", pysparkSqlFunctions.avg("col_2").over(window_2h))

Что мне нужно, так это возможность добавлять все 4 столбца (или многие другие) за один вызов, надеясь, что данные будут проходить только один раз для повышения производительности.

Answer 1

Я предпочитаю импортировать библиотеку функций как F, так как она выглядит аккуратнее, и это стандартный псевдоним, используемый в официальной документации Spark.

Звездная строка, '*', должна захватывать все текущиестолбцы в кадре данных.В качестве альтернативы вы можете заменить звездочку на *df.columns.Здесь звезда разбивает список на отдельные параметры для метода выбора.

from pyspark.sql import functions as F

df = df.select(
    "*",
    F.avg("col_1").over(window_1h).alias("col1_1h"),
    F.avg("col_1").over(window_2h).alias("col1_2h"),
    F.avg("col_2").over(window_1h).alias("col2_1h"),
    F.avg("col_2").over(window_1h).alias("col2_1h"),
)