PySpark：将DataFrame列中的列表转换为固定长度子列表

方法一：自定义UDF实现 #

逻辑直观，适合快速落地需求：

1. 创建示例DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType

spark = SparkSession.builder.appName("ListGrouping").getOrCreate()

data = [([1,2,3,4,5,6,7,8],), ([9,10,11,12,13],)]
df = spark.createDataFrame(data, ["original_list"])
df.show(truncate=False)

2. 定义分组函数并转为UDF：

def group_list(lst, group_size=4):
    return [lst[i:i+group_size] for i in range(0, len(lst), group_size)]

group_list_udf = udf(group_list, ArrayType(ArrayType(IntegerType())))

3. 应用UDF生成分组后的新列：

result_df = df.withColumn("grouped_list", group_list_udf(df.original_list))
result_df.show(truncate=False)

方法二：使用Spark内置函数（无UDF，性能更优） #

避免UDF的序列化开销，用Spark原生函数实现：

from pyspark.sql.functions import expr

result_df = df.withColumn(
    "grouped_list",
    expr("transform(sequence(0, size(original_list)-1, 4), i -> slice(original_list, i+1, 4))")
)
result_df.show(truncate=False)

补充说明 #

• 两种方法都会自动处理列表长度不是4的倍数的情况，比如列表[9,10,11,12,13]会被拆分为[[9,10,11,12], [13]]
• 若要调整分组大小，只需修改方法一中的group_size参数，或方法二中表达式里的4即可

内容的提问来源于stack exchange，提问作者chaitanya kumar Dondapati