将数组传递给Python Spark Lit函数

Question

假设我有一个numpy数组a，它包含数字1-10。所以a是[1 2 3 4 5 6 7 8 9 10]。

现在，我还有一个Python Spark数据框，我想把numpy数组a添加到其中，我想一列文字可以完成这项工作。所以我做了如下工作：

df = df.withColumn("NewColumn", F.lit(a))

这不管用。错误是“不支持的文字类型类java.util.ArrayList”。

现在，如果我只尝试数组中的一个元素，如下所示，它就可以工作了。

df = df.withColumn("NewColumn", F.lit(a[0]))

我有办法做我想做的事吗？我已经做了好几天我想完成的任务，这是我最接近完成的。我已经看了所有相关的堆栈溢出问题，但我没有得到完全的答案，我正在寻找。如有任何帮助，我们将不胜感激。谢谢。

Answer 1

数组内置函数中的for循环

您可以使用array内置函数作为

a = [1,2,3,4,5,6,7,8,9,10]
df = spark.createDataFrame([['a b c d e f g h i j '],], ['col1'])
df = df.withColumn("NewColumn", F.array([F.lit(x) for x in a]))
df.show(truncate=False)

你应该得到

+--------------------+-------------------------------+
|col1                |NewColumn                      |
+--------------------+-------------------------------+
|a b c d e f g h i j |[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]|
+--------------------+-------------------------------+
root
 |-- col1: string (nullable = true)
 |-- NewColumn: array (nullable = false)
 |    |-- element: integer (containsNull = false)

使用udf函数

#udf function
def arrayUdf():
    return a
callArrayUdf = F.udf(arrayUdf, T.ArrayType(T.IntegerType()))

#calling udf function
df = df.withColumn("NewColumn", callArrayUdf())

输出与for循环方式相同

已更新

我在粘贴下面@pault的评论

You can hide the loop using map: df.withColumn("NewColumn", F.array(map(F.lit, a)))