在PySpark中计算数组的滚动和并另存为dict？

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import MapType, StringType

spark = SparkSession.builder.appName("Array Rolling").getOrCreate()

data = [("A", [1, 2, 3, 4, 5]),
        ("B", [6, 7, 8, 9, 10]),
        ("C", [11, 12, 13, 14, 15])]
df = spark.createDataFrame(data, ["id", "values"])

def rolling_array(arr):
    result = {}
    for i in range(len(arr)):
        result[i] = sum(arr[:i+1])
    return result

rolling_udf = udf(rolling_array, MapType(StringType(), StringType()))

result_df = df.withColumn("rolling_dict", rolling_udf(df["values"]))
result_df.show(truncate=False)

+---+---------------+--------------------------------------------------+
|id |values         |rolling_dict                                      |
+---+---------------+--------------------------------------------------+
|A  |[1, 2, 3, 4, 5]|{0 -> 1, 1 -> 3, 2 -> 6, 3 -> 10, 4 -> 15}         |
|B  |[6, 7, 8, 9, 10]|{0 -> 6, 1 -> 13, 2 -> 21, 3 -> 30, 4 -> 40}       |
|C  |[11, 12, 13, 14, 15]|{0 -> 11, 1 -> 23, 2 -> 36, 3 -> 50, 4 -> 65}  |
+---+---------------+--------------------------------------------------+