PySpark-按获取分组中每个列表的大小

我有一个巨大的pyspark数据框架。我需要分组 Person 然后 collect 他们的 Budget 项目,以执行进一步的计算。 例如,

a = [('Bob', 562,"Food", "12 May 2018"), ('Bob',880,"Food","01 June 2018"), ('Bob',380,'Household'," 16 June 2018"),  ('Sue',85,'Household'," 16 July 2018"), ('Sue',963,'Household'," 16 Sept 2018")]
df = spark.createDataFrame(a, ["Person", "Amount","Budget", "Date"])

分组依据:

import pyspark.sql.functions as F
df_grouped = df.groupby('person').agg(F.collect_list("Budget").alias("data"))

架构:

root
 |-- person: string (nullable = true)
 |-- data: array (nullable = true)
 |    |-- element: string (containsNull = true)

然而,当我尝试对每个人应用UDF时,我会遇到一个内存错误。如何获取每个列表的大小(兆字节或千兆字节)( data )对于每个人?

我已经做了以下工作,但我正在 nulls

import sys
size_list_udf = F.udf(lambda data: sys.getsizeof(data)/1000, DoubleType())
df_grouped = df_grouped.withColumn("size",size_list_udf("data") )
df_grouped.show()

输出:

+------+--------------------+----+
|person|                data|size|
+------+--------------------+----+
|   Sue|[Household, House...|null|
|   Bob|[Food, Food, Hous...|null|
+------+--------------------+----+