PySpark如何不使用UDF将数组列按指定最大长度拆分为多个小数组

完全可以不用自定义UDF实现该需求，Spark 2.4及以上版本的内置数组函数即可完成，性能比自定义UDF更优，不需要额外序列化开销。

实现思路 #

核心通过内置函数生成拆分索引，再按索引切片完成拆分：

• 用size(arrayCol)计算原数组总长度
• 用sequence(0, size(arrayCol)-1, max_size)生成从0开始、步长为max_size的拆分起始索引序列
• 用transform()遍历索引序列，对每个索引调用slice(arrayCol, 索引+1, max_size)切片（Spark的slice函数下标从1开始计数，因此索引要+1），切片长度不足max_size时会自动返回剩余所有元素

代码示例 #

PySpark实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

# 初始化SparkSession
spark = SparkSession.builder.appName("array_split").getOrCreate()

# 构造测试数据
df = spark.createDataFrame([([1,2,3,4,5],)], ["arrayCol"])
print("原数据：")
df.show()

# 按max_size=2拆分数组
max_size = 2
df = df.withColumn("arrayCol", expr(f"""
    transform(
        sequence(0, size(arrayCol)-1, {max_size}),
        i -> slice(arrayCol, i+1, {max_size})
    )
"""))

print("拆分后结果：")
df.show(truncate=False)

Spark SQL实现 #

SELECT 
    transform(
        sequence(0, size(arrayCol)-1, 2),
        i -> slice(arrayCol, i+1, 2)
    ) as arrayCol
FROM 表名

输出结果 #

原数据：
+---------------+
|       arrayCol|
+---------------+
|[1, 2, 3, 4, 5]|
+---------------+

拆分后结果：
+---------------------+
|arrayCol             |
+---------------------+
|[[1, 2], [3, 4], [5]]|
+---------------------+

内容的提问来源于stack exchange，提问作者gael