通过重复 pyspark 中数组中的数字来创建列

Question

我有一个 pyspark 数据框，包含 30 行和一个包含 6 个元素的数组。

假设数组是 [5,4,3,4,1,0]。 我需要创建一个重复这 6 个数字 5 次的列。也就是说，它创建一个包含元素 [5, 4, 3, 4, 1, 0, 5, 4, 3, 4, 1, 0, 5, 4, 3, 4, 1, 0, ...] 的列并将其与初始数据框绑定。

重复函数没有帮助，因为它将整个数组重复为新数组。它创建 [5,4,3,4,1,0], [5,4,3,4,1,0], ...

如何创建此专栏？

Answer 1

从 pyspark 2.4.0 开始，您可以使用

array_repeat

和

flatten

的组合来获得所需的结果：

import pyspark.sql.functions as F

df = df.withColumn('array_repeated', F.flatten(F.array_repeat('array', 5)))

---

示例

df = spark.createDataFrame([
  ([5,4,3,4,1,0], ),
], ['array'])

df = df.withColumn('array_repeated', F.flatten(F.array_repeat('array', 5)))

df.show(truncate=False)
+------------------+------------------------------------------------------------------------------------------+
|array             |array_repeated                                                                            |
+------------------+------------------------------------------------------------------------------------------+
|[5, 4, 3, 4, 1, 0]|[5, 4, 3, 4, 1, 0, 5, 4, 3, 4, 1, 0, 5, 4, 3, 4, 1, 0, 5, 4, 3, 4, 1, 0, 5, 4, 3, 4, 1, 0]|
+------------------+------------------------------------------------------------------------------------------+