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：

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]|
# +-----+------+

Answer 2

如果可能，最好避免使用 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]|
+-----+---------+

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

2 个答案: