如何通过多次重复“另一列的大小”值来创建数组列？

Question

我想添加一个新列score，它是一个数组，其长度等于另一列values的大小，并且包含所有值2

使用列的size时出错，但如果用硬编码的数字替换它，工作正常

数据

columns = ["id","values"]
data = [("sample1", [12.0,10.0]), ("sample2", [1.0,2.0,3.0,4.0])]
rdd = spark.sparkContext.parallelize(data)

源数据帧

+-------+--------------------+
|     id|              values|
+-------+--------------------+
|sample1|        [12.0, 10.0]|
|sample2|[1.0, 2.0, 3.0, 4.0]|
+-------+--------------------+

预期产出

+-------+--------------------+--------------------+
|     id|              values|               score|
+-------+--------------------+--------------------+
|sample1|        [12.0, 10.0]|          [2, 2]    |
|sample2|[1.0, 2.0, 3.0, 4.0]|        [2, 2, 2, 2]|
+-------+--------------------+--------------------+

代码

from pyspark.sql.functions import *
df.withColumn("score",array([lit(x) for x in [2]*(size(col("values")))])).show()

低于错误值

: java.lang.RuntimeException: Unsupported literal type class java.util.ArrayList [2]

Answer 1

函数^{}仅适用于Spark 2.4+。对于旧版本，可以使用UDF完成此操作：

from pyspark.sql.functions import udf, size, lit
from pyspark.sql.types import ArrayType, IntegerType


array_repeat_udf = udf(lambda v, n: [v for _ in range(n)], ArrayType(IntegerType()))

df1 = df.withColumn('score', array_repeat_udf(lit(2), size("values")))
df1.show()

#+   -+          +      +
#|     id|              values|       score|
#+   -+          +      +
#|sample1|        [12.0, 10.0]|      [2, 2]|
#|sample2|[1.0, 2.0, 3.0, 4.0]|[2, 2, 2, 2]|
#+   -+          +      +

Answer 2

不能将Python列表与Spark列相乘。您可以使用array_repeat函数

import pyspark.sql.functions as F

df2 = df.withColumn('score', F.expr('array_repeat(2, size(values))'))
df2.show()
+   -+          +      +
|     id|              values|       score|
+   -+          +      +
|sample1|        [12.0, 10.0]|      [2, 2]|
|sample2|[1.0, 2.0, 3.0, 4.0]|[2, 2, 2, 2]|
+   -+          +      +