来自 Pyspark ArrayType 列的随机样本

我在 Pyspark 数据框中有一列，其结构如下

Column1
[a,b,c,d,e]
[c,b,d,f,g,h,i,p,l,m]

我想返回另一列，其中随机选择每行中的每个数组，以及函数中指定的数量。

所以像data.withColumn("sample", SOME_FUNCTION("column1", 5))返回：

sample
[a,b,c,d,e]
[c,b,h,i,p]

希望避免使用 python UDF，感觉应该有一个可用的函数？

这有效：

import random
def random_sample(population):
    return(random.sample(population, 5))

udf_random = F.udf(random_sample, T.ArrayType(T.StringType()))
df.withColumn("sample", udf_random("column1")).show()

但正如我所说，最好避免使用 UDF。

对于 Spark 2.4+，请使用shuffle https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.functions.shuffle and slice https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.functions.slice:

df = spark.createDataFrame([(list('abcde'),),(list('cbdfghiplm'),)],['column1'])

df.selectExpr('slice(shuffle(column1),1,5)').show()
+-----------------------------+
|slice(shuffle(column1), 1, 5)|
+-----------------------------+
|              [b, a, e, d, c]|
|              [h, f, d, l, m]|
+-----------------------------+