python - 跨数据框的pyspark窗口函数（使用lambda？）

Question

我有形状的数据框（5,100）

Idx   col_1     col_2     col_3     col_4     col_5                 
0   0.118126  0.248718  0.717174  0.551654  0.832513
1   0.823785  0.432577  0.280622  0.779573  0.917391
2   0.897924  0.388673  0.746616  0.560313  0.793524
3   0.771598  0.620364  0.406896  0.194094  0.652458
4   0.837725  0.335048  0.867820  0.493559  0.697139
5   0.721675  0.893441  0.855906  0.458607  0.806025
6   0.338725  0.732343  0.551936  0.840650  0.405240
7   0.314990  0.062280  0.145607  0.202780  0.999613
8   0.195877  0.176967  0.934697  0.028677  0.300465
9   0.643504  0.703770  0.173629  0.880687  0.290086
10  0.356232  0.400717  0.782214  0.984976  0.098428

我已将窗口定义如下：

window = Window.orderBy("Idx").rowsBetween( Window.unboundedPreceding,10)

并注册了一个自定义函数，如下所示：

@pandas_udf(FloatType(), PandasUDFType.GROUPED_AGG)
def custom_func(x: pd.Series) -> float:
    do something
    return y

我知道如何在一列上使用这个 custom_func：

df.withColumn("results", custom_func(func.col("col_1")).over(window)).orderBy("Idx")

但是，我想在整个数据框（所有 5 列）上使用这个 custom_func，并希望输出为 5 列。任何指向正确方向的指针都会非常有帮助。

标签： pythonpyspark