Apache PySpark教程：基于整数列生成对应长度的数组列

嗨，刚好Spark有个现成的函数完美匹配你的需求，不用自己写复杂的UDF，直接用array_repeat就能轻松搞定！

核心思路 #

array_repeat函数的作用就是把指定元素重复指定次数，返回一个数组。你的场景里，要重复的元素就是每行的myInt值，重复的次数也是这个值，直接传参就能得到想要的结果。

完整代码示例 #

首先确保导入必要的模块：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType
from pyspark.sql.functions import array_repeat

创建你给出的初始DataFrame：

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

# 创建目标初始DataFrame
x = spark.createDataFrame(
    [(1,), (2,), (3,)],
    StructType([StructField("myInt", IntegerType(), True)]))

新增myArr列：

# 使用array_repeat生成对应长度的数组
result_df = x.withColumn("myArr", array_repeat("myInt", "myInt"))

查看最终结果：

result_df.show()

运行后输出和你期望的完全一致：

+-----+---------+
|myInt|    myArr|
+-----+---------+
|    1|      [1]|
|    2|   [2, 2]|
|    3|[3, 3, 3]|
+-----+---------+

额外扩展 #

如果后续需求有变化，比如数组元素不需要和myInt相同，只要长度对应，也可以灵活调整：

• 生成从1到myInt的有序序列：sequence(lit(1), "myInt")，会得到[1]、[1,2]、[1,2,3]这类数组
• 生成全0数组：array_repeat(lit(0), "myInt")

不过就你给出的示例场景，array_repeat是最简洁高效的方案啦！

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