Apache pyspark 如何使用包含 n 个元素的数组创建列

Question

我有一个包含 1 列整数类型的数据框。

我想创建一个包含 n 个元素的数组的新列（n 是第一列的 #）

例如：

x = spark.createDataFrame([(1,), (2,),],StructType([ StructField("myInt", IntegerType(), True)])) 

+-----+
|myInt|
+-----+
|    1|
|    2|
|    3|
+-----+

我需要生成的数据框如下所示：

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

注意，数组中的值实际上并不重要，重要的是计数。

如果生成的数据框看起来像这样就好了：

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

Answer 1

如果可能，最好避免使用 UDF，因为它们的效率较低。您可以改用array_repeat。

import pyspark.sql.functions as F

x.withColumn('myArr', F.array_repeat(F.col('myInt'), F.col('myInt'))).show()

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

Answer 2

使用udf:

from pyspark.sql.functions import *

@udf("array<int>")
def rep_(x):
    return [x for _ in range(x)]

x.withColumn("myArr", rep_("myInt")).show()
# +-----+------+
# |myInt| myArr|
# +-----+------+
# |    1|   [1]|
# |    2|[2, 2]|
# +-----+------+